from image_prediction.label_mapper.mappers.numeric import IndexMapper
from image_prediction.label_mapper.mappers.probability import ProbabilityMapper


def test_index_label_mapper(batch_of_expected_numeric_labels, batch_of_expected_string_labels, classes):
    mapper = IndexMapper(classes)
    assert list(mapper(batch_of_expected_numeric_labels)) == batch_of_expected_string_labels


def test_array_label_mapper(
        batch_of_expected_probability_arrays, batch_of_expected_label_to_probability_mappings, classes
):
    mapper = ProbabilityMapper(classes)
    assert list(mapper(batch_of_expected_probability_arrays)) == batch_of_expected_label_to_probability_mappings