from autorag import embedding_models
from autorag.evaluation.util import cast_metrics


def test_cast_metrics():
    metric1 = ['bleu', 'meteor', 'rouge']
    metric_names, metric_params = cast_metrics(metric1)
    assert metric_names == ['bleu', 'meteor', 'rouge']
    assert metric_params == [{}, {}, {}]

    metric2 = [{'metric_name': 'bleu'}, {'metric_name': 'meteor'}, {'metric_name': 'rouge'}]
    metric_names, metric_params = cast_metrics(metric2)
    assert metric_names == ['bleu', 'meteor', 'rouge']
    assert metric_params == [{}, {}, {}]

    metric3 = [{'metric_name': 'bleu'}, {'metric_name': 'sem_score', 'embedding_model': 'openai'}]
    metric_names, metric_params = cast_metrics(metric3)
    assert metric_names == ['bleu', 'sem_score']
    assert metric_params == [{}, {'embedding_model': embedding_models['openai']}]

    metric4 = [{'metric_name': 'bleu', 'extra_param': 'extra'},
               {'metric_name': 'sem_score', 'embedding_model': 'openai', 'extra_param': 'extra'}]
    metric_names, metric_params = cast_metrics(metric4)
    assert metric_names == ['bleu', 'sem_score']
    assert metric_params == [{'extra_param': 'extra'},
                             {'embedding_model': embedding_models['openai'], 'extra_param': 'extra'}]
