import pytest

from inspect_evals.class_eval.class_eval import (
    class_eval as class_eval_task,
)
from inspect_evals.class_eval.class_eval import (
    find_code,
    record_to_sample,
)
from inspect_evals.class_eval.utils import construct_prompt

from .test_data import prompt, raw_code, record, sample, sample_code


def test_record_to_sample():
    assert record_to_sample(record) == sample


def test_find_code():
    assert find_code(raw_code) == sample_code


@pytest.mark.huggingface
def test_task():
    task = class_eval_task(few_shot=5)
    assert task.dataset.name == "FudanSELab/ClassEval"
    assert task.epochs == 5
    assert task.sandbox.type == "docker"


def test_construct_prompt():
    assert construct_prompt(record) == prompt
