Restore CSP enum basket compatibility

csp_gateway/server/gateway/csp/channels.py

@@ -2,7 +2,7 @@
 from collections import defaultdict, deque
 from contextlib import contextmanager
 from datetime import datetime
-from enum import Enum
+from enum import Enum as PyEnum
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -14,7 +14,7 @@
 )
 
 import csp
-from csp import ts
+from csp import Enum as CspEnum, ts
 from csp.impl.genericpushadapter import GenericPushAdapter
 from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from csp.impl.types.tstype import TsType, isTsType
@@ -81,7 +81,7 @@ class _SnapshotModelBaseClass(BaseModel):
 
 def _recursive_remove_enums(vals_dict):
     for k, v in list(vals_dict.items()):
-        is_enum_key = isinstance(k, Enum)
+        is_enum_key = isinstance(k, (PyEnum, CspEnum))
         is_dict_value = isinstance(v, dict)
         if is_enum_key:
             v = vals_dict.pop(k)
@@ -400,7 +400,7 @@ def _keys_for_channel(self, field: str) -> dict[Any, _NONE_TYPE] | None:
         if is_dict_basket(tstype):
             # get type of key in basket
             basket_key_type = get_dict_basket_key_type(tstype)
-            if issubclass(basket_key_type, Enum):
+            if issubclass(basket_key_type, (PyEnum, CspEnum)):
                 return {e: None for e in basket_key_type}
             else:
                 return self._dynamic_keys.get(field, {})
@@ -1081,7 +1081,7 @@ def add_send_channel(self, field: str, indexer: str | int | None = None) -> None
             tstype = tstype.typ
             self._send_channels[field, indexer] = GenericPushAdapter(tstype, name=f"manual_{field}")
 
-    def _add_send_channel_dict_basket(self, field: str, keys: list[str] | Enum) -> None:
+    def _add_send_channel_dict_basket(self, field: str, keys: list[str] | type[PyEnum] | type[CspEnum]) -> None:
         # NOTE: Do not call this directly, it is used in the factory finalization
 
         # get type of edge

csp_gateway/testing/harness.py

@@ -191,7 +191,12 @@ class GatewayAssertAttrUnsetEvent(BaseGatewayTestEvent):
     attr: str
 
     def apply(self, now, values, tick_counts, *args, **kwargs):
-        assert hasattr(values[self.channel], self.attr) == (not self.unset)
+        value = values[self.channel]
+        if hasattr(value, "_implicitly_unset"):
+            is_set = self.attr not in value._implicitly_unset() and getattr(value, self.attr, None) is not None
+        else:
+            is_set = hasattr(value, self.attr)
+        assert is_set == (not self.unset)
 
 
 class GatewayAssertLenEvent(BaseGatewayTestEvent):
@@ -256,7 +261,12 @@ class GatewayAssertIdxAttrUnsetEvent(BaseGatewayTestEvent):
     attr: str
 
     def apply(self, now, values, tick_counts, *args, **kwargs):
-        assert hasattr(values[self.channel][self.idx], self.attr) == (not self.unset)
+        value = values[self.channel][self.idx]
+        if hasattr(value, "_implicitly_unset"):
+            is_set = self.attr not in value._implicitly_unset() and getattr(value, self.attr, None) is not None
+        else:
+            is_set = hasattr(value, self.attr)
+        assert is_set == (not self.unset)
 
 
 class GatewayAssertTickedEvents(BaseGatewayTestEvent):

csp_gateway/tests/server/gateway/csp/test_channels.py

@@ -66,6 +66,12 @@ class MyGatewayChannels(GatewayChannels):
     my_array_channel: ts[Numpy1DArray[float]] = None
 
 
+def test_csp_enum_dict_basket_has_static_keys():
+    channels = MyGatewayChannels()
+
+    assert set(channels._keys_for_channel("my_enum_basket")) == set(MyEnum)
+
+
 class DerivedChannels(MyGatewayChannels):
     pass
 

csp_gateway/tests/test_harness.py

@@ -8,6 +8,16 @@
 from csp_gateway.testing.shared_helpful_classes import MyGateway, MyGatewayChannels, MyStruct
 
 
+def test_assert_attr_unset_uses_gateway_struct_field_presence():
+    h = GatewayTestHarness(test_channels=[MyGatewayChannels.my_channel])
+    h.send(MyGatewayChannels.my_channel, MyStruct())
+    h.assert_attr_unset(MyGatewayChannels.my_channel, "foo")
+    h.assert_attr_unset(MyGatewayChannels.my_channel, "my_flag", unset=False)
+
+    gateway = MyGateway(modules=[h], channels=MyGatewayChannels())
+    csp.run(gateway.graph, starttime=datetime(2020, 1, 1), endtime=timedelta(0))
+
+
 @pytest.mark.parametrize("make_invalid", (True, False))
 def test_delay(make_invalid):
     channels = [