Merge pull request #7 from SNEWS2/fix-coincidence-tier-messages

Fix coincidence tier messages

snews/models/messages.py

@@ -9,7 +9,8 @@
 # Third-party modules
 import numpy as np
 from pydantic import (BaseModel, ConfigDict, Field, NonNegativeFloat,
-                      ValidationError, field_validator, model_validator)
+                      NonNegativeInt, ValidationError, field_validator,
+                      model_validator)
 
 # Local modules
 from ..__version__ import schema_version
@@ -232,20 +233,20 @@ class RetractionMessage(DetectorMessageBase):
 
     model_config = ConfigDict(validate_assignment=True)
 
-    retract_message_uid: Optional[str] = Field(
+    retract_message_uuid: Optional[str] = Field(
         default=None,
         title="Unique message ID",
         description="Unique identifier for the message to retract"
     )
 
-    retract_latest: bool = Field(
-        default=False,
+    retract_latest_n: NonNegativeInt = Field(
+        default=0,
         title="Retract Latest Flag",
         description="True if the latest message is being retracted",
     )
 
-    retraction_reason: str = Field(
-        ...,
+    retraction_reason: Optional[str] = Field(
+        default=None,
         title="Retraction reason",
         description="Reason for retraction",
     )
@@ -257,11 +258,11 @@ def _set_tier(cls, values):
 
     @model_validator(mode="after")
     def _validate_model(self):
-        if self.retract_latest and self.retract_message_uid is not None:
-            raise ValueError("retract_message_uuid cannot be specified when retract_latest=True")
+        if self.retract_latest_n > 0 and self.retract_message_uuid is not None:
+            raise ValueError("retract_message_uuid cannot be specified when retract_latest_n > 0")
 
-        if not self.retract_latest and self.retract_message_uid is None:
-            raise ValueError("Must specify either retract_message_uuid or retract_latest=True")
+        if self.retract_latest_n == 0 and self.retract_message_uuid is None:
+            raise ValueError("Must specify either retract_message_uuid or retract_latest_n > 0")
         return self
 
 
@@ -379,7 +380,7 @@ def _set_tier(cls, values):
         values['tier'] = Tier.COINCIDENCE_TIER
         return values
 
-    @field_validator("neutrino_time_utc")
+    @field_validator("neutrino_time_utc", mode="before")
     def _validate_neutrino_time_format(cls, v: str):
         return convert_timestamp_to_ns_precision(v)
 
@@ -403,7 +404,7 @@ def _validate_neutrino_time(self):
 
 
 # .................................................................................................
-def compatible_message_types(**kwargs) -> list:
+def compatible_message_types(include_heartbeats=False, **kwargs) -> list:
     """
     Return a list of message types that are compatible with the given keyword arguments.
     """
@@ -421,6 +422,10 @@ def compatible_message_types(**kwargs) -> list:
         try:
             message_type(**kwargs)
             compatible_message_types.append(message_type)
+
+            if include_heartbeats and message_type == CoincidenceTierMessage:
+                compatible_message_types.append(HeartbeatMessage)
+
         except ValidationError:
             pass
 
@@ -435,6 +440,12 @@ def create_messages(**kwargs) -> list:
 
     messages = []
     for message_type in compatible_message_types(**kwargs):
-        messages.append(message_type(**kwargs))
+        if message_type == HeartbeatMessage:
+            messages.append(message_type(detector_status="ON", **kwargs))
+        else:
+            messages.append(message_type(**kwargs))
+
+    if len(messages) == 0:
+        raise ValueError("No compatible message types found")
 
     return messages

snews/models/timing.py

@@ -64,7 +64,6 @@ def _validate_and_cast_timestamp(cls, v):
             v = v.replace(tzinfo=None)
 
         if not isinstance(v, np.datetime64):
-            print(v)
             v = np.datetime64(v)
 
         return v

snews/schema/RetractionMessage.schema.json

@@ -140,7 +140,7 @@
       "title": "Detector Name",
       "type": "string"
     },
-    "retract_message_uid": {
+    "retract_message_uuid": {
       "anyOf": [
         {
           "type": "string"
@@ -153,22 +153,30 @@
       "description": "Unique identifier for the message to retract",
       "title": "Unique message ID"
     },
-    "retract_latest": {
-      "default": false,
+    "retract_latest_n": {
+      "default": 0,
       "description": "True if the latest message is being retracted",
+      "minimum": 0,
       "title": "Retract Latest Flag",
-      "type": "boolean"
+      "type": "integer"
     },
     "retraction_reason": {
+      "anyOf": [
+        {
+          "type": "string"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
       "description": "Reason for retraction",
-      "title": "Retraction reason",
-      "type": "string"
+      "title": "Retraction reason"
     }
   },
   "required": [
     "tier",
-    "detector_name",
-    "retraction_reason"
+    "detector_name"
   ],
   "title": "RetractionMessage",
   "type": "object"

test/models/test_message_models.py

@@ -64,7 +64,7 @@
 # Retraction message
 strategy_required_fields_retraction = {
     **strategy_required_fields_base,
-    "retract_message_uid": st.uuids(version=4).map(lambda x: str(x)),
+    "retract_message_uuid": st.uuids(version=4).map(lambda x: str(x)),
     "retraction_reason": st.text(min_size=1),
 }
 
@@ -165,10 +165,10 @@ def test_snews_message_model_retraction_required(**kwargs):
 def test_snews_message_model_retraction_validation_both_indicators(**kwargs):
     with pytest.raises(ValueError) as exc_info:
         msg = RetractionMessage(**kwargs)
-        msg.retract_latest = True
+        msg.retract_latest_n = 3
         msg.retract_message_uid = "1234567890"
 
-    assert "retract_message_uuid cannot be specified when retract_latest=True" in str(
+    assert "retract_message_uuid cannot be specified when retract_latest_n > 0" in str(
         exc_info.value
     )
 
@@ -177,9 +177,9 @@ def test_snews_message_model_retraction_validation_both_indicators(**kwargs):
 def test_snews_message_model_retraction_validation_neither_indicator(**kwargs):
     with pytest.raises(ValueError) as exc_info:
         msg = RetractionMessage(**kwargs)
-        msg.retract_latest = False
-        msg.retract_message_uid = None
+        msg.retract_latest_n = 0
+        msg.retract_message_uuid = None
 
-    assert "Must specify either retract_message_uuid or retract_latest=True" in str(
+    assert "Must specify either retract_message_uuid or retract_latest_n > 0" in str(
         exc_info.value
     )

test/unit/test_models.py

@@ -15,3 +15,36 @@ def test_import_detector_models():
     expected = models.detectors.__all__
     actual = models.detectors.__dir__()
     assert expected == actual
+
+
+# .................................................................................................
+def test_create_messages_function():
+    inputs = {
+        "detector_name": "Super-K",
+        "neutrino_time_utc": "2012-06-09T15:31:08.109876",
+        "machine_time_utc": "2012-06-09T15:30:00.009876",
+        "is_firedrill": False,
+        "is_test": True,
+    }
+
+    msgs = models.messages.create_messages(**inputs)
+    message_types = [type(msg) for msg in msgs]
+    assert message_types == [models.messages.CoincidenceTierMessage]
+
+
+# .................................................................................................
+def test_create_messages_function_with_heartbeats():
+    inputs = {
+        "detector_name": "Super-K",
+        "neutrino_time_utc": "2012-06-09T15:31:08.109876",
+        "machine_time_utc": "2012-06-09T15:30:00.009876",
+        "is_firedrill": False,
+        "is_test": True,
+    }
+
+    msgs = models.messages.create_messages(**inputs, include_heartbeats=True)
+    message_types = [type(msg) for msg in msgs]
+    assert message_types == [
+        models.messages.CoincidenceTierMessage,
+        models.messages.HeartbeatMessage
+    ]