Fix/separate union for S2 message elements

src/s2python/message.py

@@ -70,61 +70,65 @@
 )
 
 S2Message = Union[
-    FRBCActuatorDescription,
+    DDBCAverageDemandRateForecast,
+    DDBCInstruction,
+    DDBCSystemDescription,
+    DDBCTimerStatus,
     FRBCActuatorStatus,
     FRBCFillLevelTargetProfile,
-    FRBCFillLevelTargetProfileElement,
     FRBCInstruction,
     FRBCLeakageBehaviour,
-    FRBCLeakageBehaviourElement,
-    FRBCOperationMode,
-    FRBCOperationModeElement,
-    FRBCStorageDescription,
     FRBCStorageStatus,
     FRBCSystemDescription,
     FRBCTimerStatus,
     FRBCUsageForecast,
-    FRBCUsageForecastElement,
+    PEBCPowerConstraints,
     PPBCEndInterruptionInstruction,
     PPBCPowerProfileDefinition,
-    PPBCPowerSequenceContainer,
-    PPBCPowerSequence,
     PPBCPowerProfileStatus,
-    PPBCPowerSequenceContainerStatus,
-    PPBCPowerSequenceElement,
     PPBCScheduleInstruction,
     PPBCStartInterruptionInstruction,
-    PEBCAllowedLimitRange,
+    ResourceManagerDetails,
+    RevokeObject,
+    SelectControlType,
+    SessionRequest,
+    DDBCActuatorStatus,
+    FRBCInstruction,
     PEBCEnergyConstraint,
     PEBCInstruction,
-    PEBCPowerConstraints,
-    PEBCPowerEnvelope,
-    PEBCPowerEnvelopeElement,
+    Handshake,
+    HandshakeResponse,
+    InstructionStatusUpdate,
+    PowerForecast,
+    PowerMeasurement,
+    ReceptionStatus,
+]
+
+S2MessageElement = Union[
     DDBCActuatorDescription,
-    DDBCActuatorStatus,
-    DDBCAverageDemandRateForecast,
     DDBCAverageDemandRateForecastElement,
-    DDBCInstruction,
     DDBCOperationMode,
-    DDBCSystemDescription,
-    DDBCTimerStatus,
+    FRBCActuatorDescription,
+    FRBCFillLevelTargetProfileElement,
+    FRBCLeakageBehaviourElement,
+    FRBCOperationMode,
+    FRBCOperationModeElement,
+    FRBCStorageDescription,
+    FRBCUsageForecastElement,
+    PEBCAllowedLimitRange,
+    PEBCPowerEnvelope,
+    PEBCPowerEnvelopeElement,
+    PPBCPowerSequenceContainer,
+    PPBCPowerSequence,
+    PPBCPowerSequenceContainerStatus,
+    PPBCPowerSequenceElement,
     Duration,
-    Handshake,
-    HandshakeResponse,
-    InstructionStatusUpdate,
     NumberRange,
-    PowerForecast,
     PowerForecastElement,
     PowerForecastValue,
-    PowerMeasurement,
     PowerRange,
     PowerValue,
-    ReceptionStatus,
-    ResourceManagerDetails,
-    RevokeObject,
     Role,
-    SelectControlType,
-    SessionRequest,
     Timer,
     Transition,
 ]

tests/unit/message_test.py

@@ -3,6 +3,7 @@
 import importlib
 import inspect
 import pkgutil
+from typing import get_args
 
 from s2python import message
 from s2python.validate_values_mixin import S2MessageComponent
@@ -37,6 +38,14 @@ def _test_import_s2_messages(self, module_name):
             assert hasattr(
                 message, _class.__name__
             ), f"{_class} should be importable from s2_python.message"
+            if "message_id" in _class.model_fields or "subject_message_id" in _class.model_fields:
+                assert _class in get_args(message.S2Message), (
+                    f"{_class} should be typed as a s2_python.message.S2Message",
+                )
+            else:
+                assert _class in get_args(message.S2MessageElement), (
+                    f"{_class} should be typed as a s2_python.message.S2MessageElement",
+                )
 
     def test_import_s2_messages__common(self):
         self._test_import_s2_messages("s2python.common")