tests/test_validation.py

from pydantic import ValidationError
from utils.validation import generate_classification_model

def test_classification_output_validation():
    # Dynamically generate classification model
    ClassificationOutput = generate_classification_model(["Positive", "Negative"])

    # Test valid input
    valid_output = ClassificationOutput(label="Positive")
    assert valid_output.label == "Positive", "Label should be 'Positive'"

    # Test invalid input
    try:
        ClassificationOutput(label="InvalidLabel")
    except ValidationError as e:
        error_message = str(e)
        assert "Input should be 'Positive' or 'Negative'" in error_message, "Should raise validation error with correct message"