Fix cases not handled in previous supplemental attribute PRs

src/component.jl

@@ -307,14 +307,43 @@
     if !haskey(attribute_container, T)
         attribute_container[T] = Dict{Base.UUID, T}()
     end
-    attribute_container[T][get_uuid(attribute)] = attribute
-    @debug "SupplementalAttribute type $T with UUID $(get_uuid(attribute)) stored in component $(summary(component))" _group =
+
+    uuid = get_uuid(attribute)
+    if haskey(attribute_container[T], uuid)
+        throw(
+            ArgumentError(
+                "Supplemental attribute $uuid is already attached to $(summary(component))",
+            ),
+        )
+    end
+    attribute_container[T][uuid] = attribute
+    @debug "SupplementalAttribute type $T with UUID $uuid) stored in component $(summary(component))" _group =
         LOG_GROUP_SYSTEM
     return
 end
 
 """
-Return true if the component has attributes.
+Return true if the component has supplemental attributes of the given type.
+"""
+function has_supplemental_attributes(
+    ::Type{T},
+    component::InfrastructureSystemsComponent,
+) where {T <: InfrastructureSystemsSupplementalAttribute}
+    supplemental_attributes = get_supplemental_attributes_container(component)
+    if !isconcretetype(T)
+        for (k, v) in supplemental_attributes
+            if !isempty(v) && k <: T
+                return true
+            end
+        end
+    end
+    supplemental_attributes = get_supplemental_attributes_container(component)
+    !haskey(supplemental_attributes, T) && return false
+    return !isempty(supplemental_attributes[T])
+end
+
+"""
+Return true if the component has supplemental attributes.
 """
 function has_supplemental_attributes(component::InfrastructureSystemsComponent)
     container = get_supplemental_attributes_container(component)

src/supplemental_attribute.jl

@@ -24,6 +24,13 @@ function detach_component!(
     return
 end
 
+"""
+Return true if the attribute is attached to at least one component.
+"""
+function is_attached_to_component(attribute::InfrastructureSystemsSupplementalAttribute)
+    return !isempty(get_component_uuids(attribute))
+end
+
 """
 Return true if the attribute has time series data.
 """

src/supplemental_attributes.jl

@@ -76,26 +76,6 @@
     return
 end
 
-"""
-Check to see if supplemental_attribute exists.
-"""
-function has_supplemental_attributes(
-    ::Type{T},
-    component::InfrastructureSystemsComponent,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
-    supplemental_attributes = get_supplemental_attributes_container(component)
-    if !isconcretetype(T)
-        for (k, v) in supplemental_attributes
-            if !isempty(v) && k <: T
-                return true
-            end
-        end
-    end
-    supplemental_attributes = get_supplemental_attributes_container(component)
-    !haskey(supplemental_attributes, T) && return false
-    return !isempty(supplemental_attributes[T])
-end
-
 """
 Iterates over all supplemental_attributes.
 
@@ -149,6 +129,7 @@
     if isempty(supplemental_attributes.data[T])
         pop!(supplemental_attributes.data, T)
     end
+    clear_time_series_storage!(supplemental_attribute)
     return
 end
 
@@ -226,6 +207,17 @@
     return iter
 end
 
+function get_supplemental_attribute(attributes::SupplementalAttributes, uuid::Base.UUID)
+    for attr_dict in values(attributes.data)
+        attribute = get(attr_dict, uuid, nothing)
+        if !isnothing(attribute)
+            return attribute
+        end
+    end
+
+    throw(ArgumentError("No attribute with UUID=$uuid is stored"))
+end
+
 function serialize(attributes::SupplementalAttributes)
     return [serialize(y) for x in values(attributes.data) for y in values(x)]
 end

src/system_data.jl

@@ -22,6 +22,8 @@
 mutable struct SystemData <: InfrastructureSystemsType
     components::Components
     masked_components::Components
+    "Contains all attached component UUIDs, regular and masked."
+    component_uuids::Dict{Base.UUID, <:InfrastructureSystemsComponent}
     attributes::SupplementalAttributes
     time_series_params::TimeSeriesParameters
     time_series_storage::TimeSeriesStorage
@@ -69,6 +71,7 @@
     return SystemData(
         components,
         masked_components,
+        Dict{Base.UUID, InfrastructureSystemsComponent}(),
         attributes,
         TimeSeriesParameters(),
         ts_storage,
@@ -89,6 +92,7 @@
     return SystemData(
         components,
         masked_components,
+        Dict{Base.UUID, InfrastructureSystemsComponent}(),
         attributes,
         time_series_params,
         time_series_storage,
@@ -178,22 +182,22 @@
 # Arguments
 
   - `data::SystemData`: SystemData
-  - `component::InfrastructureSystemsComponent`: will store the time series reference
+  - `attribute::InfrastructureSystemsSupplementalAttribute`: will store the time series reference
   - `time_series::TimeSeriesData`: Any object of subtype TimeSeriesData
 
-Throws ArgumentError if the component is not stored in the system.
+Throws ArgumentError if the attribute is not stored in the system.
 """
 function add_time_series!(
     data::SystemData,
-    component::InfrastructureSystemsSupplementalAttribute,
+    attribute::InfrastructureSystemsSupplementalAttribute,
     time_series::TimeSeriesData;
     skip_if_present = false,
 )
     metadata_type = time_series_data_to_metadata(typeof(time_series))
     ts_metadata = metadata_type(time_series)
     attach_time_series_and_serialize!(
         data,
-        component,
+        attribute,
         ts_metadata,
         time_series;
         skip_if_present = skip_if_present,
@@ -327,6 +331,11 @@
 function compare_values(x::SystemData, y::SystemData; compare_uuids = false)
     match = true
     for name in fieldnames(SystemData)
+        if name == :component_uuids
+            # These are not serialized. They get rebuilt when the parent package adds
+            # the components.
+            continue
+        end
         val_x = getfield(x, name)
         val_y = getfield(y, name)
         if name == :time_series_storage && typeof(val_x) != typeof(val_y)
@@ -348,15 +357,34 @@
 end
 
 function remove_component!(::Type{T}, data::SystemData, name) where {T}
-    return remove_component!(T, data.components, name)
+    component = remove_component!(T, data.components, name)
+    _handle_component_removal!(data, component)
+    return component
 end
 
 function remove_component!(data::SystemData, component)
-    return remove_component!(data.components, component)
+    component = remove_component!(data.components, component)
+    _handle_component_removal!(data, component)
+    return component
 end
 
 function remove_components!(::Type{T}, data::SystemData) where {T}
-    return remove_components!(T, data.components)
+    components = remove_components!(T, data.components)
+    for component in components
+        _handle_component_removal!(data, component)
+    end
+
+    return components
+end
+
+function _handle_component_removal!(data::SystemData, component)
+    uuid = get_uuid(component)
+    if !haskey(data.component_uuids, uuid)
+        error("Bug: component = $(summary(component)) did not have its uuid stored $uuid")
+    end
+
+    pop!(data.component_uuids, uuid)
+    return
 end
 
 """
@@ -491,6 +519,10 @@
     return first_initial_timestamp, first_len
 end
 
+"""
+Provides counts of time series including attachments to components and supplemental
+attributes.
+"""
 struct TimeSeriesCounts
     components_with_time_series::Int
     supplemental_attributes_with_time_series::Int
@@ -764,30 +796,43 @@
 
 # Redirect functions to Components and TimeSeriesContainer
 
-add_component!(data::SystemData, component; kwargs...) =
+function add_component!(data::SystemData, component; kwargs...)
+    _check_duplicate_component_uuid(data, component)
     add_component!(data.components, component; kwargs...)
+    data.component_uuids[get_uuid(component)] = component
+    return
+end
 
-add_masked_component!(data::SystemData, component; kwargs...) = add_component!(
-    data.masked_components,
-    component;
-    allow_existing_time_series = true,
-    kwargs...,
-)
+function add_masked_component!(data::SystemData, component; kwargs...)
+    add_component!(
+        data.masked_components,
+        component;
+        allow_existing_time_series = true,
+        kwargs...,
+    )
+    data.component_uuids[get_uuid(component)] = component
+    return
+end
+
+function _check_duplicate_component_uuid(data::SystemData, component)
+    uuid = get_uuid(component)
+    if haskey(data.component_uuids, uuid)
+        throw(ArgumentError("Component $(summary(component)) uuid=$uuid is already stored"))
+    end
+end
 
 iterate_components(data::SystemData) = iterate_components(data.components)
 
 get_component(::Type{T}, data::SystemData, args...) where {T} =
     get_component(T, data.components, args...)
 
 function get_component(data::SystemData, uuid::Base.UUID)
-    for component in get_components(InfrastructureSystemsComponent, data)
-        if get_uuid(component) == uuid
-            return component
-        end
+    component = get(data.component_uuids, uuid, nothing)
+    if isnothing(component)
+        throw(ArgumentError("No component with uuid = $uuid is stored."))
     end
 
-    @error "no component with UUID $uuid is stored"
-    return nothing
+    return component
 end
 
 function get_components(filter_func::Function, ::Type{T}, data::SystemData) where {T}
@@ -908,6 +953,10 @@
     return get_supplemental_attributes(T, data.attributes)
 end
 
+function get_supplemental_attribute(data::SystemData, uuid::Base.UUID)
+    return get_supplemental_attributes(data.attributes, uuid)
+end
+
 function iterate_supplemental_attributes(data::SystemData)
     return iterate_supplemental_attributes(data.attributes)
 end
@@ -919,8 +968,9 @@
 )
     detach_component!(attribute, component)
     detach_supplemental_attribute!(component, attribute)
-    clear_time_series_storage!(attribute)
-    remove_supplemental_attribute!(data.attributes, attribute)
+    if !is_attached_to_component(attribute)
+        remove_supplemental_attribute!(data.attributes, attribute)
+    end
     return
 end
 
@@ -934,7 +984,6 @@
         detach_component!(attribute, component)
         detach_supplemental_attribute!(component, attribute)
     end
-    clear_time_series_storage!(attribute)
     return remove_supplemental_attribute!(data.attributes, attribute)
 end
 
@@ -950,7 +999,6 @@
             detach_supplemental_attribute!(component, attribute)
         end
         remove_supplemental_attribute!(data.attributes, attribute)
-        clear_time_series_storage!(attribute)
     end
     return
 end

test/test_supplemental_attributes.jl

@@ -47,6 +47,21 @@ end
     @test isempty(IS.get_component_uuids(geo_supplemental_attribute))
 end
 
+@testset "Test supplemental attribute attached to multiple components" begin
+    data = IS.SystemData()
+    geo_supplemental_attribute = IS.GeographicInfo()
+    component1 = IS.TestComponent("component1", 5)
+    component2 = IS.TestComponent("component2", 7)
+    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
+
+    IS.remove_supplemental_attribute!(data, component1, geo_supplemental_attribute)
+    @test IS.get_num_supplemental_attributes(data.attributes) == 1
+    IS.remove_supplemental_attribute!(data, component2, geo_supplemental_attribute)
+    @test IS.get_num_supplemental_attributes(data.attributes) == 0
+end
+
 @testset "Test iterate_SupplementalAttributes" begin
     container = IS.SupplementalAttributes(IS.InMemoryTimeSeriesStorage())
     geo_supplemental_attribute = IS.GeographicInfo()

test/test_system_data.jl

@@ -30,6 +30,7 @@
     IS.remove_component!(data, collect(components)[1])
     components = IS.get_components(IS.TestComponent, data)
     @test length(components) == 0
+    @test isempty(data.component_uuids)
 
     IS.add_component!(data, component)
     components = IS.get_components_by_name(IS.InfrastructureSystemsComponent, data, name)
@@ -228,6 +229,7 @@ end
     IS.check_components(data, IS.TestComponent)
     component = IS.get_component(IS.TestComponent, data, "component_3")
     IS.check_component(data, component)
+    @test component === IS.get_component(data, IS.get_uuid(component))
 end
 
 @testset "Test component and time series counts" begin
@@ -272,6 +274,7 @@ end
     end
 
     for c in IS.get_components(IS.TestComponent, data)
+        @test IS.has_supplemental_attributes(c)
         @test IS.has_supplemental_attributes(IS.GeographicInfo, c)
     end
 

test/test_utils.jl

@@ -50,6 +50,8 @@ struct FakeTimeSeries <: InfrastructureSystems.TimeSeriesData end
 Base.length(::FakeTimeSeries) = 42
 
 @testset "Test TimeSeriesData printing" begin
-    @test sprint(show, MIME("text/plain"), FakeTimeSeries()) ==
-          "FakeTimeSeries time_series (42):"
+    @test occursin(
+        "FakeTimeSeries time_series (42)",
+        sprint(show, MIME("text/plain"), FakeTimeSeries()),
+    )
 end