tests/test_reflection.py

import unittest

from PIL import Image

from unittest import mock

from approach.reflection import (
    build_reflection_prompt,
    make_successful_detection_request,
    openai_compatible_advisor,
    run_reflection_loop,
)


class ReflectionTest(unittest.TestCase):
    def test_pii5_crop_uses_bbox_without_resize_or_flip(self):
        image = Image.new("RGB", (12, 8), "black")
        for x in range(2, 7):
            for y in range(1, 5):
                image.putpixel((x, y), (x, y, 100))

        request = make_successful_detection_request(
            image, {"bbox": [2, 1, 5, 4], "category_name": "button"}
        )

        self.assertEqual(request.crop_image.size, (5, 4))
        self.assertEqual(request.crop_image.getpixel((0, 0)), image.getpixel((2, 1)))
        self.assertEqual(request.crop_image.getpixel((4, 0)), image.getpixel((6, 1)))
        self.assertIs(request.original_image, image)

    def test_fractional_bbox_crop_contains_the_full_detected_region(self):
        request = make_successful_detection_request(
            Image.new("RGB", (10, 10), "white"),
            {"bbox": [1.2, 2.2, 2.2, 1.2], "category_name": "button"},
        )

        self.assertEqual(request.crop_image.size, (3, 2))

    def test_reflection_loop_stops_when_advisor_has_no_concerns(self):
        calls = []

        def advisor(**kwargs):
            calls.append(kwargs)
            return {"verified": [0], "needs_refinement": [], "feedback": []}

        result = run_reflection_loop(
            Image.new("RGB", (10, 10), "white"),
            [{"bbox": [1, 1, 3, 3], "category_name": "menu"}],
            miner=lambda trace: ["menu"],
            detector=lambda candidates, previous: previous,
            advisor=advisor,
            max_iterations=10,
        )

        self.assertEqual(len(calls), 1)
        self.assertFalse(result["max_iterations_reached"])
        self.assertEqual(result["trace"][0]["verified"], [0])
        self.assertEqual(len(calls[0]["verification_requests"]), 1)
        self.assertIn("PII.6", calls[0]["pii6_prompt"])

    def test_reflection_loop_marks_max_iterations(self):
        def advisor(**kwargs):
            return {"verified": [], "needs_refinement": [0], "feedback": ["retry"]}

        result = run_reflection_loop(
            Image.new("RGB", (10, 10), "white"),
            [{"bbox": [1, 1, 3, 3], "category_name": "menu"}],
            miner=lambda trace: ["menu"],
            detector=lambda candidates, previous: previous,
            advisor=advisor,
            max_iterations=3,
        )

        self.assertTrue(result["max_iterations_reached"])
        self.assertEqual(len(result["trace"]), 3)

    def test_reflection_removes_rejected_false_positive_before_redetection(self):
        detector_calls = []

        def detector(candidates, retained):
            detector_calls.append((candidates, retained))
            return retained

        advice = [
            {"verified": [1], "needs_refinement": [0], "feedback": ["missing lever"]},
            {"verified": [0], "needs_refinement": [], "feedback": []},
        ]
        detections = [
            {"bbox": [1, 1, 2, 2], "category_name": "decoration"},
            {"bbox": [5, 5, 2, 2], "category_name": "button"},
        ]

        result = run_reflection_loop(
            Image.new("RGB", (10, 10), "white"),
            detections,
            miner=lambda trace: trace[-1]["feedback"],
            detector=detector,
            advisor=lambda **kwargs: advice.pop(0),
        )

        self.assertEqual(detector_calls[0][0], ["missing lever"])
        self.assertEqual(detector_calls[0][1], [detections[1]])
        self.assertEqual(result["detections"], [detections[1]])

    def test_openai_compatible_advisor_sends_annotated_scene_original_and_crop(self):
        fake_client = mock.Mock()
        fake_client.complete_json.return_value = {
            "verified": [0],
            "needs_refinement": [],
            "feedback": [],
        }

        with mock.patch("approach.reflection.OpenAICompatibleChatClient", return_value=fake_client):
            advisor = openai_compatible_advisor("default")
            request = make_successful_detection_request(
                Image.new("RGB", (10, 10), "white"),
                {"bbox": [1, 1, 3, 3], "category_name": "button"},
            )
            result = advisor(
                verification_requests=[request],
                unsuccessful_detection_image=Image.new("RGB", (10, 10), "white"),
                pii5_prompt="PII.5",
                pii6_prompt="PII.6",
                pii7_prompt="PII.7",
            )

        self.assertEqual(result["verified"], [0])
        call = fake_client.complete_json.call_args
        self.assertIn("PII.5", call.args[0])
        self.assertEqual(len(call.args[1]), 3)
        self.assertTrue(all(image.media_type == "image/png" for image in call.args[1]))
        self.assertIn("response_format", call.kwargs)

    def test_reflection_prompt_includes_detection_metadata(self):
        request = make_successful_detection_request(
            Image.new("RGB", (10, 10), "white"),
            {"bbox": [1, 2, 3, 4], "category_name": "slider", "score": 0.9},
        )

        prompt = build_reflection_prompt([request], "PII.5", "PII.6", "PII.7")

        self.assertIn('"category_name": "slider"', prompt)
        self.assertIn("PII.6", prompt)


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