import pytest

from data_measurements import DataMeasurementSuite


@pytest.fixture
def measurements(mock_measurements):
    mock_measurements, _ = mock_measurements
    return [
        mock_measurements.TextDuplicates.value,
        mock_measurements.TextLengths.value,
        mock_measurements.LabelDistribution.value,
    ]


@pytest.fixture
def expected_results(mock_measurements):
    _, mock_results = mock_measurements
    return mock_results


@pytest.fixture
def suite(measurements, mock_load_dataset, mock_load_metric, monkeypatch):
    return DataMeasurementSuite(
        dataset="imdb",
        measurements=measurements,
        feature="text",
        label="label",
        split="train",
        tokenizer=lambda x: x,
    )


def test_measurement_suite_initialize(suite, mock_load_dataset, measurements, monkeypatch):
    mock_load_dataset.assert_called_with("imdb", split="train")

    assert len(suite.measurements) == len(measurements)


def test_measurement_suite_run(suite, measurements, expected_results, monkeypatch):
    assert suite.measurements[0].feature == "text"
    assert suite.measurements[1].feature == "text"
    assert suite.measurements[2].feature == "label"

    results = suite.run()

    assert results["text_duplicates"] == expected_results[0]
    assert results["text_lengths"] == expected_results[1]
    assert results["label_distribution"] == expected_results[2]
