Preserve Atlas order refs in beta material context

Author: iamh2o

bloom_lims/domain/beta_lab_refs.py

@@ -601,6 +601,21 @@ def _atlas_context_for_instance(self, instance) -> dict[str, Any]:
             reference_type=self.TRF_REFERENCE_TYPE,
             value_field="atlas_trf_euid",
         )
+        atlas_order_euid = self._first_reachable_reference_value(
+            instance,
+            reference_type="order_euid",
+            value_field="order_euid",
+        )
+        atlas_order_test_euid = self._first_reachable_reference_value(
+            instance,
+            reference_type="order_test_euid",
+            value_field="order_test_euid",
+        )
+        direct_order_test_euids = self._reachable_reference_values(
+            instance,
+            reference_type="order_test_euid",
+            value_field="order_test_euid",
+        )
         atlas_test_euid = self._first_reachable_reference_value(
             instance,
             reference_type=self.TEST_REFERENCE_TYPE,
@@ -628,6 +643,16 @@ def _atlas_context_for_instance(self, instance) -> dict[str, Any]:
         )
         atlas_test_euids: list[str] = []
         seen_test_euids: set[str] = set()
+        atlas_order_test_euids: list[str] = []
+        seen_order_test_euids: set[str] = set()
+        for direct_order_test_euid in direct_order_test_euids:
+            if direct_order_test_euid in seen_order_test_euids:
+                continue
+            seen_order_test_euids.add(direct_order_test_euid)
+            atlas_order_test_euids.append(direct_order_test_euid)
+        if atlas_order_test_euid and atlas_order_test_euid not in seen_order_test_euids:
+            seen_order_test_euids.add(atlas_order_test_euid)
+            atlas_order_test_euids.append(atlas_order_test_euid)
         for direct_test_euid in direct_test_euids:
             if direct_test_euid in seen_test_euids:
                 continue
@@ -656,6 +681,16 @@ def _atlas_context_for_instance(self, instance) -> dict[str, Any]:
                 reference_type=self.TEST_REFERENCE_TYPE,
                 value_field="atlas_tenant_id",
             )
+            or self._first_reachable_reference_value(
+                instance,
+                reference_type="order_euid",
+                value_field="atlas_tenant_id",
+            )
+            or self._first_reachable_reference_value(
+                instance,
+                reference_type="order_test_euid",
+                value_field="atlas_tenant_id",
+            )
             or self._first_reachable_reference_value(
                 instance,
                 reference_type=self.TESTKIT_REFERENCE_TYPE,
@@ -685,6 +720,9 @@ def _atlas_context_for_instance(self, instance) -> dict[str, Any]:
         if not fulfillment_items:
             return {
                 "atlas_tenant_id": atlas_tenant_id,
+                "atlas_order_euid": atlas_order_euid,
+                "atlas_order_test_euid": atlas_order_test_euid,
+                "atlas_order_test_euids": atlas_order_test_euids,
                 "atlas_trf_euid": atlas_trf_euid,
                 "atlas_test_euid": atlas_test_euid,
                 "atlas_test_euids": atlas_test_euids,
@@ -709,6 +747,9 @@ def _atlas_context_for_instance(self, instance) -> dict[str, Any]:
         first = fulfillment_items[0]
         return {
             "atlas_tenant_id": first["atlas_tenant_id"],
+            "atlas_order_euid": atlas_order_euid,
+            "atlas_order_test_euid": atlas_order_test_euid,
+            "atlas_order_test_euids": atlas_order_test_euids,
             "atlas_trf_euid": first["atlas_trf_euid"] or atlas_trf_euid,
             "atlas_test_euid": atlas_test_euid or first["atlas_test_euid"],
             "atlas_test_euids": atlas_test_euids,

bloom_lims/schemas/beta_lab.py

@@ -26,6 +26,12 @@
 class AtlasFulfillmentItemReference(BaseModel):
     atlas_test_euid: str
     atlas_test_fulfillment_item_euid: str
+    atlas_order_test_euid: str | None = None
+
+
+class AtlasFulfillmentSlotReference(BaseModel):
+    atlas_order_test_euid: str
+    atlas_fulfillment_slot_euid: str
 
 
 class AtlasCollectionEventSnapshot(BaseModel):
@@ -42,6 +48,9 @@ class AtlasCollectionEventSnapshot(BaseModel):
 
 class AtlasFulfillmentContext(BaseModel):
     atlas_tenant_id: str
+    atlas_order_euid: str | None = None
+    atlas_order_test_euid: str | None = None
+    atlas_order_test_euids: list[str] = Field(default_factory=list)
     atlas_trf_euid: str | None = None
     atlas_test_euid: str | None = None
     atlas_test_euids: list[str] = Field(default_factory=list)
@@ -52,13 +61,50 @@ class AtlasFulfillmentContext(BaseModel):
     atlas_collection_event_euid: str | None = None
     collection_event_snapshot: AtlasCollectionEventSnapshot | None = None
     fulfillment_items: list[AtlasFulfillmentItemReference] = Field(default_factory=list)
+    fulfillment_slots: list[AtlasFulfillmentSlotReference] = Field(default_factory=list)
 
     @model_validator(mode="after")
     def validate_context(self) -> "AtlasFulfillmentContext":
         if not self.atlas_tenant_id.strip():
             raise ValueError("atlas_tenant_id is required")
         if self.atlas_trf_euid is not None and not self.atlas_trf_euid.strip():
             raise ValueError("atlas_trf_euid must not be empty when provided")
+        if self.atlas_order_euid is not None and not self.atlas_order_euid.strip():
+            raise ValueError("atlas_order_euid must not be empty when provided")
+        if (
+            self.atlas_order_test_euid is not None
+            and not self.atlas_order_test_euid.strip()
+        ):
+            raise ValueError("atlas_order_test_euid must not be empty when provided")
+        normalized_order_test_euids: list[str] = []
+        seen_order_test_euids: set[str] = set()
+        primary_order_test_euid = str(self.atlas_order_test_euid or "").strip()
+        if primary_order_test_euid:
+            seen_order_test_euids.add(primary_order_test_euid)
+            normalized_order_test_euids.append(primary_order_test_euid)
+        for item in self.atlas_order_test_euids:
+            clean_item = str(item or "").strip()
+            if not clean_item or clean_item in seen_order_test_euids:
+                continue
+            seen_order_test_euids.add(clean_item)
+            normalized_order_test_euids.append(clean_item)
+        for item in self.fulfillment_slots:
+            clean_item = str(item.atlas_order_test_euid or "").strip()
+            if clean_item and clean_item not in seen_order_test_euids:
+                seen_order_test_euids.add(clean_item)
+                normalized_order_test_euids.append(clean_item)
+        for item in self.fulfillment_items:
+            clean_item = str(item.atlas_order_test_euid or "").strip()
+            if clean_item and clean_item not in seen_order_test_euids:
+                seen_order_test_euids.add(clean_item)
+                normalized_order_test_euids.append(clean_item)
+        self.atlas_order_test_euids = normalized_order_test_euids
+        if self.atlas_order_test_euid is None and self.atlas_order_test_euids:
+            self.atlas_order_test_euid = self.atlas_order_test_euids[0]
+        if (self.atlas_order_test_euid or self.atlas_order_test_euids) and not self.atlas_order_euid:
+            raise ValueError(
+                "atlas_order_euid is required when atlas_order_test_euid is provided"
+            )
         if self.atlas_test_euid is not None and not self.atlas_test_euid.strip():
             raise ValueError("atlas_test_euid must not be empty when provided")
         normalized_test_euids: list[str] = []

tests/test_beta_lab_schema.py

@@ -0,0 +1,40 @@
+"""Schema regressions for Bloom beta lab payload contracts."""
+
+from __future__ import annotations
+
+import pytest
+
+from bloom_lims.schemas.beta_lab import AtlasFulfillmentContext
+
+
+def test_atlas_fulfillment_context_preserves_order_refs():
+    context = AtlasFulfillmentContext(
+        atlas_tenant_id="tenant-1",
+        atlas_order_euid="Z-AGX-ORDER1",
+        atlas_order_test_euid="Z-AGX-TEST1",
+        atlas_order_test_euids=["Z-AGX-TEST1", "Z-AGX-TEST2"],
+        fulfillment_slots=[
+            {
+                "atlas_order_test_euid": "Z-AGX-TEST2",
+                "atlas_fulfillment_slot_euid": "Z-AGX-SLOT1",
+            }
+        ],
+    )
+
+    payload = context.model_dump()
+
+    assert payload["atlas_order_euid"] == "Z-AGX-ORDER1"
+    assert payload["atlas_order_test_euid"] == "Z-AGX-TEST1"
+    assert payload["atlas_order_test_euids"] == [
+        "Z-AGX-TEST1",
+        "Z-AGX-TEST2",
+    ]
+    assert payload["fulfillment_slots"][0]["atlas_fulfillment_slot_euid"] == "Z-AGX-SLOT1"
+
+
+def test_atlas_fulfillment_context_requires_order_for_order_tests():
+    with pytest.raises(ValueError, match="atlas_order_euid is required"):
+        AtlasFulfillmentContext(
+            atlas_tenant_id="tenant-1",
+            atlas_order_test_euid="Z-AGX-TEST1",
+        )