from inspect_evals.personality.personality import record_to_sample_TRAIT


def test_record_to_sample_trait():
    """record_to_sample_TRAIT should turn a raw TRAIT record into a Sample."""
    record = {
        "question": "I see myself as someone who is talkative.",
        "response_high1": "Strongly agree",
        "response_high2": "Agree",
        "response_low1": "Disagree",
        "response_low2": "Strongly disagree",
        "personality": "Extraversion",
    }

    sample = record_to_sample_TRAIT(record)

    # input text
    assert sample.input == record["question"]
    # choice order
    assert sample.choices == [
        record["response_high1"],
        record["response_high2"],
        record["response_low1"],
        record["response_low2"],
    ]
    # fixed four-letter target list
    assert sample.target == ["A", "B", "C", "D"]
    # trait metadata
    assert sample.metadata["trait"] == "Extraversion"
