tests/test_image_utils.py

import base64
import sys
from io import BytesIO
from unittest.mock import MagicMock

import numpy as np
import pytest
from PIL import Image

sys.path.insert(0, ".")
from image_utils import (
    _b64_to_pil,
    _pil_to_cv2,
    _cv2_to_pil,
    _crop_face_thumb,
    _pil_thumb_to_b64,
    _sort_faces_left_to_right,
    _face_to_dict,
    _collect_source_pils,
)


# ── Shared fixtures ────────────────────────────────────────────────────────────


@pytest.fixture
def small_rgb_pil():
    return Image.new("RGB", (4, 4), color=(255, 0, 0))


@pytest.fixture
def small_rgb_b64(small_rgb_pil):
    buf = BytesIO()
    small_rgb_pil.save(buf, format="JPEG")
    raw = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/jpeg;base64,{raw}"


@pytest.fixture
def small_cv2_bgr():
    return np.zeros((4, 4, 3), dtype=np.uint8)


@pytest.fixture
def large_rgb_pil():
    return Image.new("RGB", (200, 200), color=(255, 0, 0))


@pytest.fixture
def fake_face():
    face = MagicMock()
    face.bbox = [40.0, 40.0, 120.0, 120.0]
    face.det_score = 0.987654
    return face


# ── _b64_to_pil ───────────────────────────────────────────────────────────────


class TestB64ToPil:
    def test_with_data_header(self, small_rgb_b64):
        pil = _b64_to_pil(small_rgb_b64)
        assert isinstance(pil, Image.Image)
        assert pil.mode == "RGB"

    def test_without_data_header(self, small_rgb_pil):
        buf = BytesIO()
        small_rgb_pil.save(buf, format="JPEG")
        raw = base64.b64encode(buf.getvalue()).decode()
        pil = _b64_to_pil(raw)
        assert isinstance(pil, Image.Image)

    def test_invalid_raises(self):
        with pytest.raises(Exception):
            _b64_to_pil("data:image/jpeg;base64,not_valid_base64!!")


# ── _pil_to_cv2 ───────────────────────────────────────────────────────────────


class TestPilToCv2:
    def test_returns_numpy_array(self, small_rgb_pil):
        arr = _pil_to_cv2(small_rgb_pil)
        assert isinstance(arr, np.ndarray)

    def test_shape_preserved(self, small_rgb_pil):
        arr = _pil_to_cv2(small_rgb_pil)
        assert arr.shape == (small_rgb_pil.height, small_rgb_pil.width, 3)

    def test_channel_order_is_bgr(self, small_rgb_pil):
        arr = _pil_to_cv2(small_rgb_pil)
        assert arr[0, 0, 2] == 255  # R in BGR
        assert arr[0, 0, 0] == 0  # B in BGR


# ── _cv2_to_pil ───────────────────────────────────────────────────────────────


class TestCv2ToPil:
    def test_returns_pil_image(self, small_cv2_bgr):
        assert isinstance(_cv2_to_pil(small_cv2_bgr), Image.Image)

    def test_mode_is_rgb(self, small_cv2_bgr):
        assert _cv2_to_pil(small_cv2_bgr).mode == "RGB"

    def test_roundtrip(self, small_rgb_pil):
        back = _cv2_to_pil(_pil_to_cv2(small_rgb_pil))
        assert back.size == small_rgb_pil.size


# ── _crop_face_thumb ──────────────────────────────────────────────────────────


class TestCropFaceThumb:
    def test_output_is_given_size(self, small_rgb_pil):
        assert _crop_face_thumb(small_rgb_pil, [1, 1, 3, 3], size=64).size == (64, 64)

    def test_bbox_near_edge_clamped(self):
        img = Image.new("RGB", (20, 20), color=(128, 128, 128))
        assert _crop_face_thumb(img, [0, 0, 5, 5], size=32).size == (32, 32)

    def test_bbox_near_right_edge_clamped(self):
        img = Image.new("RGB", (20, 20))
        assert _crop_face_thumb(img, [15, 15, 20, 20], size=32).size == (32, 32)


# ── _pil_thumb_to_b64 ─────────────────────────────────────────────────────────


class TestPilThumbToB64:
    def test_has_data_header(self, small_rgb_pil):
        assert _pil_thumb_to_b64(small_rgb_pil).startswith("data:image/jpeg;base64,")

    def test_is_decodable_base64(self, small_rgb_pil):
        _, encoded = _pil_thumb_to_b64(small_rgb_pil).split(",", 1)
        assert base64.b64decode(encoded)[:2] == b"\xff\xd8"  # JPEG magic bytes


# ── _sort_faces_left_to_right ─────────────────────────────────────────────────


class TestSortFaces:
    def _make_face(self, x1):
        f = MagicMock()
        f.bbox = [float(x1), 0.0, float(x1 + 10), 10.0]
        return f

    def test_sorted_ascending(self):
        faces = [self._make_face(50), self._make_face(10), self._make_face(30)]
        result = _sort_faces_left_to_right(faces)
        assert [f.bbox[0] for f in result] == sorted(f.bbox[0] for f in faces)

    def test_empty_list(self):
        assert _sort_faces_left_to_right([]) == []

    def test_single_face(self):
        face = self._make_face(5)
        assert _sort_faces_left_to_right([face]) == [face]


# ── _face_to_dict ─────────────────────────────────────────────────────────────


class TestFaceToDict:
    def test_has_required_keys(self, large_rgb_pil, fake_face):
        assert {"idx", "thumbnail", "bbox", "det_score"} <= _face_to_dict(
            0, fake_face, large_rgb_pil
        ).keys()

    def test_idx_matches(self, large_rgb_pil, fake_face):
        assert _face_to_dict(3, fake_face, large_rgb_pil)["idx"] == 3

    def test_bbox_is_list_of_ints(self, large_rgb_pil, fake_face):
        bbox = _face_to_dict(0, fake_face, large_rgb_pil)["bbox"]
        assert isinstance(bbox, list)
        assert all(isinstance(v, int) for v in bbox)

    def test_det_score_rounded_to_3dp(self, large_rgb_pil, fake_face):
        assert _face_to_dict(0, fake_face, large_rgb_pil)["det_score"] == round(
            fake_face.det_score, 3
        )

    def test_thumbnail_has_header(self, large_rgb_pil, fake_face):
        assert _face_to_dict(0, fake_face, large_rgb_pil)["thumbnail"].startswith(
            "data:image/jpeg;base64,"
        )


# ── _collect_source_pils ──────────────────────────────────────────────────────


class TestCollectSourcePils:
    def test_returns_pil_list(self, small_rgb_b64):
        result = _collect_source_pils([0, 1], {"0": small_rgb_b64, "1": small_rgb_b64})
        assert len(result) == 2
        assert all(isinstance(p, Image.Image) for p in result)

    def test_skips_missing_indices(self, small_rgb_b64):
        result = _collect_source_pils([0, 1], {"0": small_rgb_b64})
        assert len(result) == 1