tests/src/test_benchmarks.py

import pytest

from src.benchmarks import LongDocBenchmarks, QABenchmarks
from src.envs import BENCHMARK_VERSION_LIST


# Ref: https://github.com/AIR-Bench/AIR-Bench/blob/4b27b8a8f2047a963805fcf6fb9d74be51ec440c/docs/available_tasks.md
# 24.05
# | Task | dev | test |
# | ---- | --- | ---- |
# | Long-Doc | 4 | 11 |
# | QA | 54 | 53 |
#
# 24.04
# | Task | test |
# | ---- | ---- |
# | Long-Doc | 15 |
# | QA | 13 |

@pytest.mark.parametrize(
    "num_datasets_dict",
    [
        {
            "air_bench_2404": 13,
            "air_bench_2405": 53
        }
    ]
)
def test_qa_benchmarks(num_datasets_dict):
    assert len(QABenchmarks) == len(BENCHMARK_VERSION_LIST)
    for benchmark_list in list(QABenchmarks):
        version_slug = benchmark_list.name
        assert num_datasets_dict[version_slug] == len(benchmark_list.value)


@pytest.mark.parametrize(
    "num_datasets_dict",
    [
        {
            "air_bench_2404": 15,
            "air_bench_2405": 11
        }
    ]
)
def test_doc_benchmarks(num_datasets_dict):
    assert len(LongDocBenchmarks) == len(BENCHMARK_VERSION_LIST)
    for benchmark_list in list(LongDocBenchmarks):
        version_slug = benchmark_list.name
        assert num_datasets_dict[version_slug] == len(benchmark_list.value)