From 05bd96773322ccc411807dd7347b9a1332dbb553 Mon Sep 17 00:00:00 2001
From: GotthardG <51994228+GotthardG@users.noreply.github.com>
Date: Mon, 29 Jun 2026 16:05:05 +0200
Subject: [PATCH] adding test for rotation screening diffraction image
 ingestion logic to aaredb and service methods

---
 .../test_rotation_screening_diffraction.py    | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 tests/unit/daq/test_rotation_screening_diffraction.py

diff --git a/tests/unit/daq/test_rotation_screening_diffraction.py b/tests/unit/daq/test_rotation_screening_diffraction.py
new file mode 100644
index 00000000..d777958d
--- /dev/null
+++ b/tests/unit/daq/test_rotation_screening_diffraction.py
@@ -0,0 +1,66 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from aare.daq.operations.rotation.service import RotationService
+
+
+def _img(number, angle=0.0, spots=0, res=0.0):
+    return SimpleNamespace(number=number, angle=angle, spots=spots, res=res)
+
+
+def _make_service(jfjoch, aare):
+    ctx = SimpleNamespace(deps=SimpleNamespace(jfjoch=jfjoch, aare=aare))
+    return RotationService(context=ctx, logger=MagicMock())
+
+
+def test_screening_diffraction_uploads_one_image_per_wedge():
+    jfjoch = MagicMock()
+    jfjoch.get_diffraction_image.side_effect = lambda image_id, **kw: f"jpeg-{image_id}".encode()
+    aare = MagicMock()
+    svc = _make_service(jfjoch, aare)
+
+    scan_result = SimpleNamespace(images=[_img(0, 0.0, 245, 2.1), _img(1, 90.0, 12, 8.5)])
+    svc._ingest_screening_diffraction(SimpleNamespace(db_id=123), scan_result)
+
+    # one fetch + one upload per wedge, with spots + resolution-estimate ring on
+    assert jfjoch.get_diffraction_image.call_count == 2
+    for call in jfjoch.get_diffraction_image.call_args_list:
+        assert call.kwargs["show_spots"] is True
+        assert call.kwargs["show_res_est"] is True
+
+    assert aare.upload_jpg.call_count == 2
+    first = aare.upload_jpg.call_args_list[0]
+    assert first.args[:3] == (123, "123_screening_diffraction_0", b"jpeg-0")
+    assert "Screening diffraction (wedge 1" in first.kwargs["message"]
+    assert "245 spots" in first.kwargs["message"]
+
+
+def test_screening_diffraction_skips_images_without_id():
+    jfjoch = MagicMock()
+    jfjoch.get_diffraction_image.return_value = b"jpeg"
+    aare = MagicMock()
+    svc = _make_service(jfjoch, aare)
+
+    svc._ingest_screening_diffraction(
+        SimpleNamespace(db_id=7),
+        SimpleNamespace(images=[SimpleNamespace(number=None), _img(3)]),
+    )
+
+    jfjoch.get_diffraction_image.assert_called_once_with(3, show_spots=True, show_res_est=True)
+    assert aare.upload_jpg.call_count == 1
+
+
+def test_screening_diffraction_tolerates_fetch_failure():
+    jfjoch = MagicMock()
+    jfjoch.get_diffraction_image.side_effect = [Exception("buffer miss"), b"jpeg-6"]
+    aare = MagicMock()
+    svc = _make_service(jfjoch, aare)
+
+    svc._ingest_screening_diffraction(
+        SimpleNamespace(db_id=9),
+        SimpleNamespace(images=[_img(5), _img(6)]),
+    )
+
+    # first wedge failed and was skipped; second still uploaded
+    assert aare.upload_jpg.call_count == 1
+    assert aare.upload_jpg.call_args.args[1] == "9_screening_diffraction_6"