Use callbacks when fetching images

adit/core/processors.py

@@ -18,6 +18,7 @@
 from .types import DicomLogEntry, ProcessingResult
 from .utils.dicom_dataset import QueryDataset, ResultDataset
 from .utils.dicom_operator import DicomOperator
+from .utils.dicom_utils import write_dataset
 from .utils.sanitize import sanitize_dirname
 
 logger = logging.getLogger(__name__)
@@ -157,7 +158,7 @@ def _download_dicoms(
             pseudonym = None
             patient_folder = download_folder / sanitize_dirname(self.transfer_task.patient_id)
 
-        study = self._fetch_study()
+        study = self._find_study()
         modalities = study.ModalitiesInStudy
 
         modalities = [
@@ -173,7 +174,7 @@ def _download_dicoms(
                 # TODO: this seems to be very ineffective as we do a c-find for every series
                 # in a study and check if it is a wanted series. Better fetch all series of
                 # a study and check then.
-                for series in self._fetch_series_list(series_uid=series_uid):
+                for series in self._find_series_list(series_uid=series_uid):
                     modalities.add(series.Modality)
 
         study_date = study.StudyDate
@@ -209,7 +210,7 @@ def _download_dicoms(
 
         return patient_folder
 
-    def _fetch_study(self) -> ResultDataset:
+    def _find_study(self) -> ResultDataset:
         studies = list(
             self.source_operator.find_studies(
                 QueryDataset.create(
@@ -256,7 +257,7 @@ def _fetch_study(self) -> ResultDataset:
 
         return study
 
-    def _fetch_series_list(self, series_uid: str) -> list[ResultDataset]:
+    def _find_series_list(self, series_uid: str) -> list[ResultDataset]:
         series_list = list(
             self.source_operator.find_series(
                 QueryDataset.create(
@@ -292,22 +293,32 @@ def _download_study(
         modifier: Callable,
         series_uids: list[str] | None = None,
     ) -> None:
+        def callback(ds: Dataset | None) -> None:
+            if ds is None:
+                return
+
+            modifier(ds)
+
+            folder_path = Path(study_folder)
+            file_name = f"{ds.SOPInstanceUID}.dcm"
+            file_path = folder_path / file_name
+            folder_path.mkdir(parents=True, exist_ok=True)
+            write_dataset(ds, file_path)
+
         if series_uids:
             for series_uid in series_uids:
-                self.source_operator.download_series(
+                self.source_operator.fetch_series(
                     patient_id=patient_id,
                     study_uid=study_uid,
                     series_uid=series_uid,
-                    dest_folder=study_folder,
-                    modifier=modifier,
+                    callback=callback,
                 )
         else:
             # If no series are explicitly chosen then download all series of the study
-            self.source_operator.download_study(
+            self.source_operator.fetch_study(
                 patient_id=patient_id,
                 study_uid=study_uid,
-                dest_folder=study_folder,
-                modifier=modifier,
+                callback=callback,
             )
 
     def _modify_dataset(

adit/core/tests/test_processors.py

@@ -77,12 +77,7 @@ def test_transfer_to_server_succeeds(
     result = processor.process()
 
     # Assert
-    source_operator_mock.download_study.assert_called_with(
-        task.patient_id,
-        task.study_uid,
-        mocker.ANY,
-        modifier=mocker.ANY,
-    )
+    source_operator_mock.fetch_study.assert_called_with(task.patient_id, task.study_uid, mocker.ANY)
 
     upload_path = dest_operator_mock.upload_instances.call_args.args[0]
     assert upload_path.match(f"*/{study.PatientID}")
@@ -133,10 +128,7 @@ def test_transfer_to_folder_succeeds(
     result = processor.process()
 
     # Assert
-    download_path = source_operator_mock.download_study.call_args.kwargs["dest_folder"]
-    assert download_path.match(
-        rf"adit_{job._meta.app_label}_{job.id}_20200101_kai/1001/20190923-080000-CT"
-    )
+    source_operator_mock.fetch_study.assert_called_with(task.patient_id, task.study_uid, mocker.ANY)
 
     assert result["status"] == TransferTask.Status.SUCCESS
     assert result["message"] == "Transfer task completed successfully."

adit/core/tests/utils/test_dicom_operator.py

@@ -1,8 +1,6 @@
 import asyncio
-import filecmp
 import threading
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from time import sleep
 
 from django.conf import settings
@@ -133,14 +131,18 @@ def test_download_series_with_c_get(
     association.send_c_get.return_value = dicom_test_helper.create_successful_c_get_response()
     path = Path(settings.BASE_DIR) / "samples" / "dicoms"
     ds = read_dataset(next(path.rglob("*.dcm")))
-    with TemporaryDirectory() as tmp_dir:
-        # Act
-        dicom_operator.download_series(
-            ds.PatientID, ds.StudyInstanceUID, ds.SeriesInstanceUID, Path(tmp_dir)
-        )
+    received_ds = []
 
-        # Assert
-        association.send_c_get.assert_called_once()
+    # Act
+    dicom_operator.fetch_series(
+        ds.PatientID, ds.StudyInstanceUID, ds.SeriesInstanceUID, lambda ds: received_ds.append(ds)
+    )
+
+    # Assert
+    association.send_c_get.assert_called_once()
+
+    # TODO: This test could be improved, unfortunately the callback of  fetch_series will never get
+    # called when we just mock send_c_get. And so we can't assert anything on received_ds.
 
 
 def test_download_series_with_c_move(
@@ -186,15 +188,16 @@ async def on_subscribe(topic: str):
     # Make sure transmit server is started
     sleep(0.5)
 
-    with TemporaryDirectory() as tmp_dir:
-        # Act
-        dicom_operator.download_series(
-            ds.PatientID, ds.StudyInstanceUID, ds.SeriesInstanceUID, Path(tmp_dir)
-        )
-
-        # Assert
-        assert subscribed_topic == (
-            f"{dicom_operator.server.ae_title}\\{ds.StudyInstanceUID}\\{ds.SeriesInstanceUID}"
-        )
-        association.send_c_move.assert_called_once()
-        assert filecmp.cmp(next(Path(tmp_dir).iterdir()), file_path)
+    received_ds = []
+
+    # Act
+    dicom_operator.fetch_series(
+        ds.PatientID, ds.StudyInstanceUID, ds.SeriesInstanceUID, lambda ds: received_ds.append(ds)
+    )
+
+    # Assert
+    assert subscribed_topic == (
+        f"{dicom_operator.server.ae_title}\\{ds.StudyInstanceUID}\\{ds.SeriesInstanceUID}"
+    )
+    association.send_c_move.assert_called_once()
+    assert received_ds[0] == ds