From b3fd7faacd5de10ad586405889c16f6e35fd5e1d Mon Sep 17 00:00:00 2001
From: Daniel Thom <daniel.thom@nrel.gov>
Date: Tue, 9 Jan 2024 09:17:19 -0700
Subject: [PATCH] Get components with a supplemental attribute

---
 src/system_data.jl                   | 14 ++++++++++++--
 test/test_internal.jl                |  2 +-
 test/test_supplemental_attributes.jl |  2 ++
 test/test_system_data.jl             | 25 +++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/system_data.jl b/src/system_data.jl
index f1338ee16..925fd5641 100644
--- a/src/system_data.jl
+++ b/src/system_data.jl
@@ -866,6 +866,11 @@ end
 get_components_by_name(::Type{T}, data::SystemData, args...) where {T} =
     get_components_by_name(T, data.components, args...)
 
+function get_components(data::SystemData, attribute::SupplementalAttribute)
+    uuids = get_component_uuids(attribute)
+    return [get_component(data, x) for x in uuids]
+end
+
 function get_masked_components(
     ::Type{T},
     data::SystemData,
@@ -950,8 +955,13 @@ end
 _get_system_basename(system_file) = splitext(basename(system_file))[1]
 _get_secondary_basename(system_basename, name) = system_basename * "_" * name
 
-add_supplemental_attribute!(data::SystemData, component, info; kwargs...) =
-    add_supplemental_attribute!(data.attributes, component, info; kwargs...)
+function add_supplemental_attribute!(data::SystemData, component, info; kwargs...)
+    if isnothing(get_component(typeof(component), data, get_name(component)))
+        throw(ArgumentError("$(summary(component)) is not attached to the system"))
+    end
+
+    return add_supplemental_attribute!(data.attributes, component, info; kwargs...)
+end
 
 function get_supplemental_attributes(
     filter_func::Function,
diff --git a/test/test_internal.jl b/test/test_internal.jl
index 2839f4c32..d62877024 100644
--- a/test/test_internal.jl
+++ b/test/test_internal.jl
@@ -1,7 +1,7 @@
 @testset "Test assign_new_uuid_internal" begin
     component = IS.TestComponent("component", 5)
     uuid1 = IS.get_uuid(component)
-    IS.assign_new_uuid_internal(component)
+    IS.assign_new_uuid_internal!(component)
     @test uuid1 != IS.get_uuid(component)
 end
 
diff --git a/test/test_supplemental_attributes.jl b/test/test_supplemental_attributes.jl
index 5dd5b53c0..e94f8843b 100644
--- a/test/test_supplemental_attributes.jl
+++ b/test/test_supplemental_attributes.jl
@@ -52,6 +52,8 @@ end
     geo_supplemental_attribute = IS.GeographicInfo()
     component1 = IS.TestComponent("component1", 5)
     component2 = IS.TestComponent("component2", 7)
+    IS.add_component!(data, component1)
+    IS.add_component!(data, component2)
     IS.add_supplemental_attribute!(data, component1, geo_supplemental_attribute)
     IS.add_supplemental_attribute!(data, component2, geo_supplemental_attribute)
     @test IS.get_num_supplemental_attributes(data.attributes) == 1
diff --git a/test/test_system_data.jl b/test/test_system_data.jl
index 944470e17..34e583665 100644
--- a/test/test_system_data.jl
+++ b/test/test_system_data.jl
@@ -311,6 +311,15 @@ end
     attr2 = IS.TestSupplemental(; value = 2.0)
     component1 = IS.TestComponent("component1", 5)
     component2 = IS.TestComponent("component2", 7)
+
+    @test_throws ArgumentError IS.add_supplemental_attribute!(
+        data,
+        component1,
+        geo_supplemental_attribute,
+    )
+
+    IS.add_component!(data, component1)
+    IS.add_component!(data, component2)
     IS.add_supplemental_attribute!(data, component1, geo_supplemental_attribute)
     IS.add_supplemental_attribute!(data, component2, geo_supplemental_attribute)
     IS.add_supplemental_attribute!(data, component1, attr1)
@@ -394,3 +403,19 @@ end
     @test IS.get_supplemental_attribute(data, uuid3) ===
           IS.get_supplemental_attribute(component2, uuid3)
 end
+
+@testset "Test retrieval of components with a supplemental attribute" begin
+    data = IS.SystemData()
+    geo_supplemental_attribute = IS.GeographicInfo()
+    component1 = IS.TestComponent("component1", 5)
+    component2 = IS.TestComponent("component2", 7)
+    IS.add_component!(data, component1)
+    IS.add_component!(data, component2)
+    IS.add_supplemental_attribute!(data, component1, geo_supplemental_attribute)
+    IS.add_supplemental_attribute!(data, component2, geo_supplemental_attribute)
+    components = IS.get_components(data, geo_supplemental_attribute)
+    @test length(components) == 2
+    sort!(components; by = x -> x.name)
+    @test components[1] === component1
+    @test components[2] === component2
+end