test_cases = [
    {
        "predictions": [0, 1, 0, 1, 0, 1],
        "references":  [0, 1, 0, 1, 0, 1],
        "sample_weight": None,
        "adjusted": False,
        "result": {"balanced_accuracy": 0}
    },
    {
        "predictions": [0, 0, 1, 1, 1, 1],
        "references":  [0, 0, 0, 0, 1, 1],
        "sample_weight": None,
        "adjusted": False,
        "result": {"balanced_accuracy": 0}
    },
    {
        "predictions": [0, 1, 1, 0, 1, 2],
        "references":  [0, 1, 2, 0, 1, 2],
        "sample_weight": None,
        "adjusted": False,
        "result": {"balanced_accuracy": 0}
    },
    {
        "predictions": [0, 0, 1, 2, 1, 2],
        "references":  [0, 0, 0, 0, 1, 2],
        "sample_weight": None,
        "adjusted": False,
        "result": {"balanced_accuracy": 0}
    },
    {
        "predictions": [0, 1, 1, 0, 0, 1],
        "references":  [0, 1, 0, 1, 0, 1],
        "sample_weight": [0.5, 0.7, 0.8, 0.9, 1.0, 0.6],
        "adjusted": False,
        "result": {"balanced_accuracy": 0}
    },
    {
        "predictions": [0, 1, 1, 0, 0, 1],
        "references":  [0, 1, 0, 1, 0, 1],
        "sample_weight": None,
        "adjusted": True,
        "result": {"balanced_accuracy": 0}
    },
]

import pytest
from evaluate import load
from sklearn.metrics import balanced_accuracy_score

@pytest.mark.parametrize("test_case", test_cases)
def test_balanced_accuracy(test_case):
    metric = load("hyperml/balanced_accuracy")
    result = metric.compute(
        predictions=test_case["predictions"],
        references=test_case["references"],
        sample_weight=test_case["sample_weight"],
        adjusted=test_case["adjusted"]
    )
    assert result["balanced_accuracy"] == balanced_accuracy_score(y_pred=test_case["predictions"], y_true=test_case["references"], sample_weight=test_case["sample_weight"], adjusted=test_case["adjusted"])
    assert result == test_case["result"]