from aarecommon.models.models import MLBoxType, MLOutputModel


def test_add_box_generates_unique_keys():
    model = MLOutputModel()
    key1 = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8)
    key2 = model.add_box(MLBoxType.CRYSTAL, (5, 6, 7, 8), 0.9)

    assert key1 == "Crystal"
    assert key2 == "Crystal_2"


def test_get_best_for_class_returns_highest_confidence():
    model = MLOutputModel()
    model.add_box(MLBoxType.PIN, (1, 1, 2, 2), 0.3)
    model.add_box(MLBoxType.PIN, (3, 3, 4, 4), 0.7)

    best = model.get_best_for_class(MLBoxType.PIN)

    assert best is not None
    assert best.conf == 0.7


def test_get_best_for_class_returns_none_when_missing():
    model = MLOutputModel()
    assert model.get_best_for_class(MLBoxType.CRYSTAL) is None