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
    PYRAMID_FLOOR,
    RebuildError,
    below_pyramid_floor,
    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 PyramidFloorTests(unittest.TestCase):
    """The l09 defect: a view too small to store as the protocol's pyramid."""

    def make_source(self, height):
        root = tempfile.mkdtemp(prefix="floor-")
        image = cv2.resize(artwork(), (300, height), interpolation=cv2.INTER_NEAREST)
        cv2.imwrite(os.path.join(root, "sku_x1.jpg"), image)
        return root

    def test_a_small_view_is_reported(self):
        root = self.make_source(305)
        small = below_pyramid_floor({"sku_x1": ["sku_x1.jpg"]}, root)
        self.assertEqual(len(small), 1, small)
        self.assertEqual(small[0]["sku"], "sku_x1")
        self.assertEqual(small[0]["height"], 305)
        self.assertEqual(small[0]["floor"], PYRAMID_FLOOR)

    def test_a_full_size_view_is_not_reported(self):
        root = self.make_source(1000)
        self.assertEqual(below_pyramid_floor({"sku_x1": ["sku_x1.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}'))


class BadShelfListFailsBeforeTheSwapTests(unittest.TestCase):
    """A bad --shelf-tags must not leave a replaced catalog behind.

    The list is read before write_catalog runs, so a malformed file or a missing
    path fails with the tool's own message and the destination is untouched. Read
    after the swap, the same error would escape as a traceback over a catalog
    that had already been replaced.
    """

    def make_source(self):
        root = tempfile.mkdtemp(prefix="guard-src-")
        cv2.imwrite(os.path.join(root, "sku_x1.jpg"), artwork())
        with open(os.path.join(root, "stock.csv"), "w", encoding="utf-8") as handle:
            handle.write("SKU_ID,product_name,price,number_in,number_out,number_on\n")
            handle.write("sku_x1,Alpha,100,1,0,1\n")
        return root

    def run_tool(self, shelf_tags):
        from rebuild_catalog import main

        source = self.make_source()
        output = os.path.join(tempfile.mkdtemp(prefix="guard-out-"), "catalog")
        with self.assertRaises(SystemExit) as caught:
            main(["--source", source, "--output", output, "--csv",
                  os.path.join(source, "stock.csv"), "--min-views", "1",
                  "--ocr", "none", "--shelf-tags", shelf_tags])
        return str(caught.exception), output

    def test_a_malformed_list_fails_cleanly_and_writes_nothing(self):
        path = os.path.join(tempfile.mkdtemp(prefix="bad-tags-"), "tags.json")
        with open(path, "w", encoding="utf-8") as handle:
            handle.write('{"sku_x1": 1}')
        message, output = self.run_tool(path)
        self.assertIn("expected a JSON list", message)
        self.assertFalse(os.path.exists(output), "the catalog was replaced before the check")

    def test_a_missing_list_fails_cleanly_and_writes_nothing(self):
        message, output = self.run_tool("/tmp/definitely-not-a-shelf-list.json")
        self.assertIn("No such file", message)
        self.assertFalse(os.path.exists(output), "the catalog was replaced before the check")


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