product-recognition-demo/recognization/code/test_rebuild_catalog.py

"""
Contract tests for the enrollment tool's reporting checks.

The refusal rules (size floors, session leakage, name/OCR, duplicates) are what
the tool exists for; these tests cover the two checks that report rather than
refuse, because both are properties of the catalog as a whole and neither is
visible from a single view: cross-SKU near-duplicates, and the completeness
rule that decides whether a bay is live-ready.
"""

from __future__ import annotations

import os
import sys
import tempfile
import unittest

current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from rebuild_catalog import (  # noqa: E402
    RebuildError,
    live_ready_verdict,
    load_shelf_skus,
    text_discriminated_siblings,
)


def artwork(mark=None):
    """A package face; `mark` adds the small text that separates siblings."""
    image = np.full((400, 300, 3), 200, dtype="uint8")
    cv2.rectangle(image, (20, 20), (280, 380), (40, 40, 180), -1)
    cv2.putText(image, "MAYO", (60, 200), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (255, 255, 255), 4)
    if mark:
        cv2.putText(image, mark, (60, 300), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 3)
    return image


class SiblingFlagTests(unittest.TestCase):
    """The near-unknown case: same artwork, different product."""

    def make_source(self, images):
        root = tempfile.mkdtemp(prefix="siblings-")
        for name, image in images.items():
            cv2.imwrite(os.path.join(root, name), image)
        return root

    def test_shared_artwork_across_skus_is_flagged(self):
        root = self.make_source({"sku_x1.jpg": artwork(), "sku_x2.jpg": artwork("HALF")})
        flagged = text_discriminated_siblings({"sku_x1": ["sku_x1.jpg"], "sku_x2": ["sku_x2.jpg"]}, root)
        self.assertEqual(len(flagged), 1, flagged)
        self.assertEqual(flagged[0]["sku"], "sku_x1")
        self.assertEqual(flagged[0]["sibling_sku"], "sku_x2")
        self.assertGreater(flagged[0]["cosine"], 0.98)

    def test_distinct_artwork_is_not_flagged(self):
        other = np.full((400, 300, 3), 30, dtype="uint8")
        cv2.circle(other, (150, 200), 90, (200, 200, 40), -1)
        root = self.make_source({"sku_x1.jpg": artwork(), "sku_x2.jpg": other})
        self.assertEqual(
            text_discriminated_siblings({"sku_x1": ["sku_x1.jpg"], "sku_x2": ["sku_x2.jpg"]}, root), []
        )

    def test_two_views_of_one_sku_are_not_siblings(self):
        """The same-SKU duplicate rule already covers these; this flag is cross-SKU."""
        root = self.make_source({"sku_x1.jpg": artwork(), "sku_x1_2.jpg": artwork()})
        self.assertEqual(
            text_discriminated_siblings({"sku_x1": ["sku_x1.jpg", "sku_x1_2.jpg"]}, root), []
        )


class CompletenessVerdictTests(unittest.TestCase):
    """The enrollment protocol's completeness rule, as a verdict."""

    def test_without_a_shelf_list_the_verdict_is_unknown(self):
        """Not a guess: the rule needs the shelf's SKU list to mean anything."""
        verdict = live_ready_verdict({"sku_a": ["a.jpg"]}, {"sku_a": "A"}, None)
        self.assertIsNone(verdict["live_ready"])
        self.assertIn("shelf", verdict["reason"])

    def test_every_shelf_sku_enrolled_is_live_ready(self):
        verdict = live_ready_verdict({"sku_a": ["a.jpg"]}, {"sku_a": "A"}, {"sku_a"})
        self.assertTrue(verdict["live_ready"])

    def test_a_text_only_row_satisfies_the_rule(self):
        """A CSV row with no image still names and vetoes from its text."""
        verdict = live_ready_verdict({"sku_a": ["a.jpg"]}, {"sku_a": "A", "sku_b": "B"}, {"sku_a", "sku_b"})
        self.assertTrue(verdict["live_ready"])

    def test_a_shelf_sku_absent_from_the_catalog_fails(self):
        verdict = live_ready_verdict({"sku_a": ["a.jpg"]}, {"sku_a": "A"}, {"sku_a", "sku_z"})
        self.assertFalse(verdict["live_ready"])
        self.assertEqual(verdict["missing_from_catalog"], ["sku_z"])


class ShelfTagListTests(unittest.TestCase):
    def write(self, text):
        path = os.path.join(tempfile.mkdtemp(prefix="tags-"), "tags.json")
        with open(path, "w", encoding="utf-8") as handle:
            handle.write(text)
        return path

    def test_a_json_list_is_read(self):
        self.assertEqual(load_shelf_skus(self.write('["sku_a", "sku_b"]')), {"sku_a", "sku_b"})

    def test_one_sku_per_line_is_read(self):
        self.assertEqual(load_shelf_skus(self.write("sku_a\nsku_b\n")), {"sku_a", "sku_b"})

    def test_no_path_means_no_list(self):
        self.assertIsNone(load_shelf_skus(None))

    def test_a_json_object_is_refused(self):
        with self.assertRaises(RebuildError):
            load_shelf_skus(self.write('{"sku_a": 1}'))


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