tests/test_inference_smoke.py

import tempfile
import unittest
from pathlib import Path

from src.inference import ClassifiedPrediction, _aggregate_classification, _iter_image_paths


class TestInferenceHelpers(unittest.TestCase):
    def test_iter_image_paths_filters_and_recurses(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            (root / "a.jpg").write_bytes(b"x")
            (root / "b.txt").write_text("ignore")
            nested = root / "nested"
            nested.mkdir()
            (nested / "c.PNG").write_bytes(b"y")

            paths = _iter_image_paths(root)
            names = [p.name for p in paths]

            self.assertEqual(names, ["a.jpg", "c.PNG"])

    def test_aggregate_classification_majority_vote(self):
        classified = [
            ClassifiedPrediction(path=Path("1.jpg"), pred_idx=0, confidence=0.80),
            ClassifiedPrediction(path=Path("2.jpg"), pred_idx=0, confidence=0.60),
            ClassifiedPrediction(path=Path("3.jpg"), pred_idx=2, confidence=0.95),
        ]
        labels = ["Class A", "Class B", "Class C"]

        label, mean_conf = _aggregate_classification(classified, labels)

        self.assertEqual(label, labels[0])
        self.assertAlmostEqual(mean_conf, 0.70, places=6)

    def test_aggregate_classification_tie_break_by_mean_confidence(self):
        classified = [
            ClassifiedPrediction(path=Path("1.jpg"), pred_idx=0, confidence=0.60),
            ClassifiedPrediction(path=Path("2.jpg"), pred_idx=0, confidence=0.50),
            ClassifiedPrediction(path=Path("3.jpg"), pred_idx=1, confidence=0.90),
            ClassifiedPrediction(path=Path("4.jpg"), pred_idx=1, confidence=0.80),
        ]
        labels = ["Case 0", "Case 1"]

        label, mean_conf = _aggregate_classification(classified, labels)

        self.assertEqual(label, labels[1])
        self.assertAlmostEqual(mean_conf, 0.85, places=6)


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