Files
mlflow--mlflow/tests/entities/test_assessment_source.py
2026-07-13 13:22:34 +08:00

77 lines
2.9 KiB
Python

import pytest
from mlflow.entities.assessment_source import AssessmentSource
from mlflow.exceptions import MlflowException
def test_assessment_source_equality():
source1 = AssessmentSource(source_type="HUMAN", source_id="user_1")
source2 = AssessmentSource(source_type="HUMAN", source_id="user_1")
source3 = AssessmentSource(source_type="LLM_JUDGE", source_id="ai_1")
source5 = AssessmentSource(source_type="HUMAN", source_id="user_2")
assert source1 == source2 # Same type and ID
assert source1 != source3 # Different type
assert source1 != source5 # Different ID
def test_assessment_source_properties():
source = AssessmentSource(source_type="HUMAN", source_id="user_1")
assert source.source_type == "HUMAN"
assert source.source_id == "user_1"
def test_assessment_source_to_from_dictionary():
source = AssessmentSource(source_type="HUMAN", source_id="user_1")
source_dict = source.to_dictionary()
expected_dict = {
"source_type": "HUMAN",
"source_id": "user_1",
}
assert source_dict == expected_dict
recreated_source = AssessmentSource.from_dictionary(source_dict)
assert recreated_source == source
def test_assessment_source_type_validation():
# Valid source types
AssessmentSource(source_type="HUMAN", source_id="user_1")
AssessmentSource(source_type="LLM_JUDGE", source_id="judge_1")
# Deprecated source type for backward compatibility
source = AssessmentSource(source_type="AI_JUDGE", source_id="ai_1")
assert source.source_type == "LLM_JUDGE"
# Invalid source type
with pytest.raises(MlflowException, match="Invalid assessment source type"):
AssessmentSource(source_type="ROBOT", source_id="robot_1")
def test_assessment_source_case_insensitivity():
# Valid source types with different cases
source_1 = AssessmentSource(source_type="human", source_id="user_1")
source_2 = AssessmentSource(source_type="Human", source_id="user_2")
source_3 = AssessmentSource(source_type="HUMAN", source_id="user_3")
source_4 = AssessmentSource(source_type="llm_judge", source_id="judge_1")
source_5 = AssessmentSource(source_type="Llm_Judge", source_id="judge_2")
source_6 = AssessmentSource(source_type="LLM_JUDGE", source_id="judge_3")
# Verify that the source type is normalized to uppercase
assert source_1.source_type == "HUMAN"
assert source_2.source_type == "HUMAN"
assert source_3.source_type == "HUMAN"
assert source_4.source_type == "LLM_JUDGE"
assert source_5.source_type == "LLM_JUDGE"
assert source_6.source_type == "LLM_JUDGE"
def test_ai_judge_deprecation_warning():
with pytest.warns(FutureWarning, match="AI_JUDGE is deprecated. Use LLM_JUDGE instead."):
source = AssessmentSource(source_type="AI_JUDGE", source_id="ai_1")
assert source.source_type == "LLM_JUDGE"
assert source.source_id == "ai_1"