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")


class OrientationTests(unittest.TestCase):
    """A packshot whose identifying text is not horizontal.

    The やまや めんたいマヨネーズ packshot reads its brand only when rotated;
    the upright pass returns stray kana. Checking one orientation refused a
    correct image, which is the worst failure this check has - it rejects good
    data while claiming to validate it.
    """
    class FakeOcr:
        """Reads a fixed string per rotation.

        Rotations are identified by aspect ratio, not by exact shape: the
        caller upscales before OCR, so the fake sees scaled dimensions. The
        upright and 180-degree passes keep the source's orientation (wide),
        the two 90-degree passes swap it (tall). A corner marker cannot tell
        them apart, because `cv2.rotate` carries the corner value through.
        """

        name = "fake"

        def __init__(self, by_orientation):
            self.by_orientation = by_orientation

        def read(self, image):
            class Line:
                def __init__(self, text):
                    self.text = text

            class Read:
                def __init__(self, text):
                    self.lines = [Line(part) for part in text.split("|") if part]

            height, width = image.shape[:2]
            key = "wide" if width > height else "tall"
            return Read(self.by_orientation.get(key, ""))

    def marked(self, height, width):
        return np.zeros((height, width, 3), dtype="uint8")

    def test_a_rotated_read_wins_when_it_matches_the_name(self):
        """The upright pass reads noise; the rotated pass reads the product."""
        from rebuild_catalog import best_orientation_text

        engine = self.FakeOcr({"wide": "つくつ|た", "tall": "明太手屋さんが|マヨネーズ"})
        text = best_orientation_text(engine, self.marked(40, 60), "やまや めんたいマヨネーズタイプ 200g")
        self.assertIn("明太手屋さんが", text)

    def test_the_upright_read_wins_when_it_matches(self):
        from rebuild_catalog import best_orientation_text

        engine = self.FakeOcr({"wide": "ピュアセレクト|マヨネーズ|450g", "tall": "ノイズ"})
        text = best_orientation_text(engine, self.marked(40, 60), "ピュアセレクト マヨネーズ 450g")
        self.assertIn("ピュアセレクト", text)

    def test_without_a_name_the_longest_read_wins(self):
        from rebuild_catalog import best_orientation_text

        engine = self.FakeOcr({"wide": "ab", "tall": "abcdefgh"})
        self.assertEqual(best_orientation_text(engine, self.marked(40, 60)), "abcdefgh")

    def test_audit_keeps_a_view_whose_name_only_reads_rotated(self):
        """The regression: `audit` must try every orientation before refusing.

        This is the boundary that matters. `best_orientation_text` passing its
        own unit test proves nothing about enrollment - the bug was that
        `audit` called a single-orientation read, so a correct packshot whose
        brand text runs vertically was refused as a name/OCR mismatch.
        """
        from rebuild_catalog import audit

        root = tempfile.mkdtemp(prefix="orient-")
        # Wider than tall, so the upright pass is the "wide" read (noise) and
        # only a 90-degree rotation produces the "tall" read that matches.
        cv2.imwrite(os.path.join(root, "sku_x1.jpg"), self.marked(300, 400))
        engine = self.FakeOcr({"wide": "つくつ|た", "tall": "キユーピー|エッグケア|305g"})
        kept, refused = audit(root, min_views=1, names={"sku_x1": "キューピー エッグケア 305g"},
                              ocr_engine=engine)
        self.assertEqual(kept, {"sku_x1": ["sku_x1.jpg"]}, refused)
        self.assertEqual(refused, [])

    def test_audit_still_refuses_a_view_the_name_never_matches(self):
        """The check must keep its teeth: a wrong product is still refused."""
        from rebuild_catalog import audit

        root = tempfile.mkdtemp(prefix="orient-bad-")
        cv2.imwrite(os.path.join(root, "sku_x1.jpg"), self.marked(400, 300))
        engine = self.FakeOcr({"wide": "キューピー|マヨネーズ", "tall": "キューピー|マヨネーズ"})
        kept, refused = audit(root, min_views=1, names={"sku_x1": "やまや めんたいマヨネーズタイプ 200g"},
                              ocr_engine=engine)
        self.assertEqual(kept, {})
        self.assertEqual(len(refused), 1, refused)
        self.assertIn("name/OCR", refused[0]["reason"])