from datasets import Dataset

from data_measurements.measurements import LabelDistribution


def test_label_distribution_initialize(mock_load_metric):
    LabelDistribution(feature="text")
    mock_load_metric.assert_called_once_with("label_distribution")


def test_label_distribution_run(mock_load_metric):
    dataset = Dataset.from_dict(
        {
            "text": ["Hello", "World", "Hello", "Foo Bar"],
            "label": [1, 2, 1, 1],
        }
    )
    label_distribution = LabelDistribution(feature="label")
    label_distribution.measure(dataset)

    label_distribution.metric.compute.assert_called_once_with(data=dataset["label"])


def test_label_distribution_figure(mock_load_metric):
    dataset = Dataset.from_dict(
        {
            "text": ["Hello", "World", "Hello", "Foo Bar"],
            "label": [1, 2, 1, 1],
        }
    )

    # TODO: Some kind of assertion?
    label_distribution = LabelDistribution(feature="label")
    results = label_distribution.measure(dataset)

    results.to_figure()
