Make MWLStorage#update_status smarter

Such that it only performs valid status transitions

Author: carlosmartinez

src/services/dicom/c_store.py

@@ -8,6 +8,7 @@
 from services.dicom.image_compressor import ImageCompressor
 from services.dicom.validation_failure_notifier import ValidationFailureNotifier
 from services.dicom.validator import DicomValidationError, DicomValidator
+from services.mwl import MWLStatus
 from services.storage import InstanceExistsError, MWLStorage, PACSStorage
 
 logger = logging.getLogger(__name__)
@@ -102,7 +103,7 @@ def _mark_in_progress(self, accession_number: str) -> None:
         if not self.mwl_storage or not accession_number:
             return
         try:
-            self.mwl_storage.mark_in_progress(accession_number)
+            self.mwl_storage.update_status(accession_number, MWLStatus.IN_PROGRESS.value)
         except Exception as e:
             logger.error(f"Failed to mark worklist item in progress: {e}", exc_info=True)
 

src/services/storage.py

@@ -7,6 +7,7 @@
 from typing import Dict, List, Optional
 
 from models import WorklistItem
+from services.mwl import MWLStatus
 
 logger = logging.getLogger(__name__)
 
@@ -284,13 +285,25 @@ class WorklistItemNotFoundError(Exception):
     pass
 
 
+class InvalidStatusTransitionError(Exception):
+    """Raised when a requested status transition is not permitted."""
+
+    pass
+
+
 class DuplicateWorklistItemError(Exception):
     """Raised when a worklist item with the same accession number already exists."""
 
     pass
 
 
 class MWLStorage(Storage):
+    _STATUS_TRANSITIONS: dict[MWLStatus, MWLStatus] = {
+        MWLStatus.IN_PROGRESS: MWLStatus.SCHEDULED,
+        MWLStatus.COMPLETED: MWLStatus.IN_PROGRESS,
+        MWLStatus.DISCONTINUED: MWLStatus.IN_PROGRESS,
+    }
+
     def __init__(self, db_path: str = "/var/lib/pacs/worklist.db"):
         """
         Initialize Worklist storage.
@@ -459,16 +472,24 @@ def update_status(
         self, accession_number: str, status: str, mpps_instance_uid: Optional[str] = None
     ) -> Optional[str]:
         """
-        Update the status of a worklist item.
+        Transition a worklist item to a new status, enforcing valid state transitions.
 
         Args:
             accession_number: The accession number to update
-            status: New status (SCHEDULED, IN PROGRESS, COMPLETED, DISCONTINUED)
+            status: Target status
             mpps_instance_uid: Optional MPPS instance UID
 
         Returns:
             source_message_id if item was updated, None if not found
+
+        Raises:
+            InvalidStatusTransitionError: If the transition is not permitted
         """
+        target = MWLStatus(status)
+        if target not in self._STATUS_TRANSITIONS:
+            raise InvalidStatusTransitionError(f"Cannot transition to '{status}'")
+        from_status = self._STATUS_TRANSITIONS[target]
+
         with self._get_connection() as conn:
             cursor = conn.execute(
                 """
@@ -477,31 +498,21 @@ def update_status(
                     mpps_instance_uid = COALESCE(?, mpps_instance_uid),
                     updated_at = CURRENT_TIMESTAMP
                 WHERE accession_number = ?
-            """,
-                (status, mpps_instance_uid, accession_number),
+                  AND status = ?
+                """,
+                (status, mpps_instance_uid, accession_number, from_status.value),
             )
             conn.commit()
 
             if cursor.rowcount == 0:
                 return None
 
             result = conn.execute(
-                "SELECT source_message_id FROM worklist_items WHERE accession_number = ?", (accession_number,)
+                "SELECT source_message_id FROM worklist_items WHERE accession_number = ?",
+                (accession_number,),
             ).fetchone()
-
             return result["source_message_id"] if result is not None else None
 
-    def mark_in_progress(self, accession_number: str) -> bool:
-        """
-        Transition worklist item from SCHEDULED to IN PROGRESS.
-
-        Returns True if the status was updated, False if not found or already past SCHEDULED.
-        """
-        item = self.get_worklist_item(accession_number)
-        if item is None or item.status != "SCHEDULED":
-            return False
-        return self.update_status(accession_number, "IN PROGRESS") is not None
-
     def update_study_instance_uid(self, accession_number: str, study_instance_uid: str) -> bool:
         """
         Update the study instance UID for a worklist item.

tests/integration/test_c_find_returns_worklist_items.py

@@ -265,8 +265,6 @@ def test_cfind_filters_by_modality(self, event, storage):
                 source_message_id="MSGID123456",
             )
         )
-        storage.update_status("ACC234567", "SCHEDULED")
-
         event.identifier.ScheduledProcedureStepSequence[0].Modality = "MG"
 
         results = list(CFind(storage).call(event))

tests/integration/test_end_to_end_relay_to_upload.py

@@ -134,9 +134,6 @@ async def test_full_flow_relay_to_upload(
         assert worklist_items[0].accession_number == TEST_ACCESSION_NUMBER
         assert worklist_items[0].patient_id == TEST_PATIENT_ID
 
-        # Update status to SCHEDULED so it appears in C-FIND results
-        mwl_storage.update_status(TEST_ACCESSION_NUMBER, "SCHEDULED")
-
         # ===== STEP 2: Query worklist via C-FIND =====
         mwl_server.start()
         try:

tests/integration/test_request_cfind_on_worklist.py

@@ -31,7 +31,6 @@ def with_pacs_server(self, tmp_dir):
                 source_message_id="MSGID123456",
             )
         )
-        storage.update_status("ACC123456", "SCHEDULED")
         storage.store_worklist_item(
             WorklistItem(
                 accession_number="ACC234567",
@@ -48,8 +47,6 @@ def with_pacs_server(self, tmp_dir):
                 source_message_id="MSGID234567",
             )
         )
-        storage.update_status("ACC234567", "SCHEDULED")
-
         server.start()
 
         yield

tests/services/dicom/test_c_store.py

@@ -124,26 +124,26 @@ def test_validation_failure_notifies_manage(self, mock_storage, mock_event):
         mock_notifier.notify.assert_called_once_with("action-uuid-123", "DICOM validation failed: Missing required tag")
         mock_mwl.get_source_message_id.assert_called_once_with("ABC123")
 
-    def test_mark_in_progress_called_on_success(self, mock_storage, mock_event):
+    def test_worklist_marked_in_progress_on_success(self, mock_storage, mock_event):
         mock_mwl = Mock(spec=MWLStorage)
         subject = CStore(mock_storage, mwl_storage=mock_mwl)
 
         assert subject.call(mock_event) == SUCCESS
 
-        mock_mwl.mark_in_progress.assert_called_once_with("ABC123")
+        mock_mwl.update_status.assert_called_once_with("ABC123", "IN PROGRESS")
 
-    def test_mark_in_progress_not_called_on_failure(self, mock_storage, mock_event):
+    def test_worklist_not_updated_on_store_failure(self, mock_storage, mock_event):
         mock_storage.store_instance.side_effect = Exception("store failed")
         mock_mwl = Mock(spec=MWLStorage)
         subject = CStore(mock_storage, mwl_storage=mock_mwl)
 
         assert subject.call(mock_event) == FAILURE
 
-        mock_mwl.mark_in_progress.assert_not_called()
+        mock_mwl.update_status.assert_not_called()
 
-    def test_mark_in_progress_error_does_not_fail_store(self, mock_storage, mock_event):
+    def test_worklist_update_error_does_not_fail_store(self, mock_storage, mock_event):
         mock_mwl = Mock(spec=MWLStorage)
-        mock_mwl.mark_in_progress.side_effect = Exception("db error")
+        mock_mwl.update_status.side_effect = Exception("db error")
         subject = CStore(mock_storage, mwl_storage=mock_mwl)
 
         assert subject.call(mock_event) == SUCCESS

tests/services/test_storage.py

@@ -6,7 +6,7 @@
 from pydicom.uid import generate_uid
 
 from models import WorklistItem
-from services.storage import MWLStorage, PACSStorage, WorklistItemNotFoundError
+from services.storage import InvalidStatusTransitionError, MWLStorage, PACSStorage, WorklistItemNotFoundError
 
 
 @pytest.fixture
@@ -228,20 +228,19 @@ def test_get_worklist_item_returns_none(self, mwl_storage):
 
     def test_update_status(self, mwl_storage, result):
         item = self._insert_item(mwl_storage, result)
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
         returned = mwl_storage.update_status(item.accession_number, "COMPLETED")
 
         assert returned == item.source_message_id
-
         assert mwl_storage.get_worklist_item(item.accession_number).status == "COMPLETED"
 
-    def test_update_status_with_no_update(self, mwl_storage):
-        result = mwl_storage.update_status("DOES_NOT_EXIST", "COMPLETED")
-
-        assert result is None
+    def test_update_status_returns_none_when_not_found(self, mwl_storage):
+        assert mwl_storage.update_status("DOES_NOT_EXIST", "IN PROGRESS") is None
 
     def test_update_status_with_mpps(self, mwl_storage, result):
         item = self._insert_item(mwl_storage, result)
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
         returned = mwl_storage.update_status(
             item.accession_number,
@@ -250,14 +249,19 @@ def test_update_status_with_mpps(self, mwl_storage, result):
         )
 
         assert returned == item.source_message_id
+        assert mwl_storage.get_worklist_item(item.accession_number).mpps_instance_uid == "some-uid"
 
-        with mwl_storage._get_connection() as conn:
-            row = conn.execute(
-                "SELECT mpps_instance_uid FROM worklist_items WHERE accession_number = ?",
-                (item.accession_number,),
-            ).fetchone()
+    def test_update_status_raises_on_invalid_target(self, mwl_storage, result):
+        item = self._insert_item(mwl_storage, result)
 
-        assert row["mpps_instance_uid"] == "some-uid"
+        with pytest.raises(InvalidStatusTransitionError):
+            mwl_storage.update_status(item.accession_number, "SCHEDULED")  # SCHEDULED is never a valid target
+
+    def test_update_status_returns_none_on_wrong_state(self, mwl_storage, result):
+        item = self._insert_item(mwl_storage, result)
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
+
+        assert mwl_storage.update_status(item.accession_number, "IN PROGRESS") is None
 
     def test_update_study_instance_uid(self, mwl_storage, result):
         item = self._insert_item(mwl_storage, result)
@@ -288,6 +292,7 @@ def test_delete_worklist_item_raises(self, mwl_storage):
     def test_mpps_instance_exists(self, mwl_storage, result):
         uid = generate_uid()
         item = self._insert_item(mwl_storage, result)
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
         mwl_storage.update_status(item.accession_number, "COMPLETED", mpps_instance_uid=uid)
 
@@ -299,6 +304,7 @@ def test_mpps_instance_not_exists(self, mwl_storage):
     def test_get_worklist_item_by_mpps_instance_uid(self, mwl_storage, result):
         uid = generate_uid()
         item = self._insert_item(mwl_storage, result)
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
         mwl_storage.update_status(item.accession_number, "COMPLETED", mpps_instance_uid=uid)
 
@@ -313,18 +319,17 @@ def test_get_worklist_item_by_mpps_instance_uid(self, mwl_storage, result):
     def test_get_worklist_item_by_mpps_instance_uid_returns_none(self, mwl_storage):
         assert mwl_storage.get_worklist_item_by_mpps_instance_uid("nope") is None
 
-    def test_mark_in_progress(self, mwl_storage, result):
-        item = self._insert_item(mwl_storage, result)  # inserted with status='SCHEDULED'
+    def test_update_status_scheduled_to_in_progress(self, mwl_storage, result):
+        item = self._insert_item(mwl_storage, result)
 
-        assert mwl_storage.mark_in_progress(item.accession_number) is True
+        mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
         assert mwl_storage.get_worklist_item(item.accession_number).status == "IN PROGRESS"
 
-    def test_mark_in_progress_is_idempotent(self, mwl_storage, result):
+    def test_update_status_in_progress_to_discontinued(self, mwl_storage, result):
         item = self._insert_item(mwl_storage, result)
         mwl_storage.update_status(item.accession_number, "IN PROGRESS")
 
-        assert mwl_storage.mark_in_progress(item.accession_number) is False
+        mwl_storage.update_status(item.accession_number, "DISCONTINUED")
 
-    def test_mark_in_progress_returns_false_when_not_found(self, mwl_storage):
-        assert mwl_storage.mark_in_progress("DOES_NOT_EXIST") is False
+        assert mwl_storage.get_worklist_item(item.accession_number).status == "DISCONTINUED"