sciagent/be0/tests/test_imagehub_segmentation.py

"""Unit tests for the ImageHub segmentation-linking domain service.

The service was built with INJECTED infrastructure (put_blob / sniff_meta / safe_name)
precisely so the domain rules can be exercised with fakes — no Postgres, no MinIO.
Covers: parent validation (bad uuid / not-found / not-an-image), the mask path
namespacing, organ-label fallback, and the empty-payload guard.
"""
from __future__ import annotations

import unittest
import uuid

from src.imagehub_segmentation import MaskUpload, SegmentationError, SegmentationService
from src.initiative_db.models import ImagehubDataset, ImagehubDatasetFile


class _FakeResult:
    def __init__(self, value):
        self._value = value

    def scalar_one_or_none(self):
        return self._value


class _FakeSession:
    """Minimal AsyncSession stand-in. The 1st execute() resolves the parent file;
    every later execute() resolves the 'existing mask at this path' lookup (None =
    create new). Records added rows."""

    def __init__(self, parent=None, existing=None):
        self._parent = parent
        self._existing = existing
        self.added: list = []
        self.exec_calls = 0
        self.flushes = 0

    async def execute(self, _stmt):
        self.exec_calls += 1
        return _FakeResult(self._parent if self.exec_calls == 1 else self._existing)

    async def get(self, _model, _key):
        return None  # blob absent → service will add it

    def add(self, obj):
        self.added.append(obj)

    async def flush(self):
        self.flushes += 1


def _safe_name(name):
    base = (name or "").strip().replace("\\", "/").rsplit("/", 1)[-1]
    return base or "file"


async def _put_blob(data, media_type):
    return {
        "sha256": "deadbeef" + str(len(data)),
        "size": len(data),
        "bucket": "imagehub-blobs",
        "key": "blobs/de/ad/deadbeef",
        "media_type": media_type or "application/octet-stream",
        "deduped": False,
    }


def _sniff(_filename, _data, _media):
    return {"format": "nifti", "shape": [4, 4, 4]}


def _service(session):
    return SegmentationService(session, put_blob=_put_blob, sniff_meta=_sniff, safe_name=_safe_name)


def _image_parent(dataset_id, logical_path="ct.nii.gz"):
    return ImagehubDatasetFile(
        id=uuid.uuid4(), dataset_id=dataset_id, logical_path=logical_path, file_kind="image"
    )


def _mask(filename="liver.nii.gz", organ="Gan", data=b"xyz"):
    return MaskUpload(filename=filename, data=data, media_type="application/gzip", organ_label=organ)


class TestMaskLogicalPath(unittest.TestCase):
    def test_namespaces_under_parent_stem(self):
        svc = _service(_FakeSession())
        self.assertEqual(svc._mask_logical_path("ct.nii.gz", "liver.nii.gz"), "ct.seg/liver.nii.gz")
        self.assertEqual(svc._mask_logical_path("scan.nii", "a.nii.gz"), "scan.seg/a.nii.gz")
        self.assertEqual(svc._mask_logical_path("study.dcm", "k.nii.gz"), "study.seg/k.nii.gz")

    def test_mask_path_cannot_collide_with_a_real_image_path(self):
        # A real file's logical_path never contains '/', a mask path always does.
        svc = _service(_FakeSession())
        self.assertIn("/", svc._mask_logical_path("ct.nii.gz", "ct.nii.gz"))


class TestLinkMasks(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.ds = ImagehubDataset(id=uuid.uuid4(), owner_user_id=uuid.uuid4())
        self.uid = uuid.uuid4()

    async def test_happy_path_links_mask_to_image(self):
        parent = _image_parent(self.ds.id)
        sess = _FakeSession(parent=parent, existing=None)
        rows = await _service(sess).link_masks(self.ds, str(parent.id), [_mask()], self.uid)
        self.assertEqual(len(rows), 1)
        r = rows[0]
        self.assertEqual(r.file_kind, "segmentation")
        self.assertEqual(r.parent_file_id, parent.id)
        self.assertEqual(r.organ_label, "Gan")
        self.assertEqual(r.logical_path, "ct.seg/liver.nii.gz")
        self.assertEqual(r.dataset_id, self.ds.id)

    async def test_organ_label_falls_back_to_filename(self):
        parent = _image_parent(self.ds.id)
        sess = _FakeSession(parent=parent)
        rows = await _service(sess).link_masks(self.ds, str(parent.id), [_mask(organ="  ")], self.uid)
        self.assertEqual(rows[0].organ_label, "liver.nii.gz")

    async def test_bad_parent_uuid_is_404(self):
        with self.assertRaises(SegmentationError) as ctx:
            await _service(_FakeSession()).link_masks(self.ds, "not-a-uuid", [_mask()], self.uid)
        self.assertEqual(ctx.exception.status, 404)

    async def test_parent_not_found_is_404(self):
        sess = _FakeSession(parent=None)
        with self.assertRaises(SegmentationError) as ctx:
            await _service(sess).link_masks(self.ds, str(uuid.uuid4()), [_mask()], self.uid)
        self.assertEqual(ctx.exception.status, 404)

    async def test_attaching_to_a_mask_is_422(self):
        mask_parent = ImagehubDatasetFile(
            id=uuid.uuid4(), dataset_id=self.ds.id, logical_path="ct.seg/x.nii.gz", file_kind="segmentation"
        )
        sess = _FakeSession(parent=mask_parent)
        with self.assertRaises(SegmentationError) as ctx:
            await _service(sess).link_masks(self.ds, str(mask_parent.id), [_mask()], self.uid)
        self.assertEqual(ctx.exception.status, 422)

    async def test_empty_payload_is_422(self):
        parent = _image_parent(self.ds.id)
        sess = _FakeSession(parent=parent)
        with self.assertRaises(SegmentationError) as ctx:
            await _service(sess).link_masks(self.ds, str(parent.id), [], self.uid)
        self.assertEqual(ctx.exception.status, 422)

    async def test_all_empty_byte_masks_is_422(self):
        parent = _image_parent(self.ds.id)
        sess = _FakeSession(parent=parent)
        with self.assertRaises(SegmentationError) as ctx:
            await _service(sess).link_masks(self.ds, str(parent.id), [_mask(data=b"")], self.uid)
        self.assertEqual(ctx.exception.status, 422)


if __name__ == "__main__":
    unittest.main()