Get components with a supplemental attribute

NREL-Sienna · Jan 9, 2024 · b3fd7fa · b3fd7fa
1 parent b871ecc
commit b3fd7fa
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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