1  import pytest
 2  
 3  from mlflow.entities.assessment_source import AssessmentSource
 4  from mlflow.exceptions import MlflowException
 5  
 6  
 7  def test_assessment_source_equality():
 8      source1 = AssessmentSource(source_type="HUMAN", source_id="user_1")
 9      source2 = AssessmentSource(source_type="HUMAN", source_id="user_1")
10      source3 = AssessmentSource(source_type="LLM_JUDGE", source_id="ai_1")
11      source5 = AssessmentSource(source_type="HUMAN", source_id="user_2")
12  
13      assert source1 == source2  # Same type and ID
14      assert source1 != source3  # Different type
15      assert source1 != source5  # Different ID
16  
17  
18  def test_assessment_source_properties():
19      source = AssessmentSource(source_type="HUMAN", source_id="user_1")
20  
21      assert source.source_type == "HUMAN"
22      assert source.source_id == "user_1"
23  
24  
25  def test_assessment_source_to_from_dictionary():
26      source = AssessmentSource(source_type="HUMAN", source_id="user_1")
27      source_dict = source.to_dictionary()
28  
29      expected_dict = {
30          "source_type": "HUMAN",
31          "source_id": "user_1",
32      }
33      assert source_dict == expected_dict
34  
35      recreated_source = AssessmentSource.from_dictionary(source_dict)
36      assert recreated_source == source
37  
38  
39  def test_assessment_source_type_validation():
40      # Valid source types
41      AssessmentSource(source_type="HUMAN", source_id="user_1")
42      AssessmentSource(source_type="LLM_JUDGE", source_id="judge_1")
43  
44      # Deprecated source type for backward compatibility
45      source = AssessmentSource(source_type="AI_JUDGE", source_id="ai_1")
46      assert source.source_type == "LLM_JUDGE"
47  
48      # Invalid source type
49      with pytest.raises(MlflowException, match="Invalid assessment source type"):
50          AssessmentSource(source_type="ROBOT", source_id="robot_1")
51  
52  
53  def test_assessment_source_case_insensitivity():
54      # Valid source types with different cases
55      source_1 = AssessmentSource(source_type="human", source_id="user_1")
56      source_2 = AssessmentSource(source_type="Human", source_id="user_2")
57      source_3 = AssessmentSource(source_type="HUMAN", source_id="user_3")
58      source_4 = AssessmentSource(source_type="llm_judge", source_id="judge_1")
59      source_5 = AssessmentSource(source_type="Llm_Judge", source_id="judge_2")
60      source_6 = AssessmentSource(source_type="LLM_JUDGE", source_id="judge_3")
61  
62      # Verify that the source type is normalized to uppercase
63      assert source_1.source_type == "HUMAN"
64      assert source_2.source_type == "HUMAN"
65      assert source_3.source_type == "HUMAN"
66      assert source_4.source_type == "LLM_JUDGE"
67      assert source_5.source_type == "LLM_JUDGE"
68      assert source_6.source_type == "LLM_JUDGE"
69  
70  
71  def test_ai_judge_deprecation_warning():
72      with pytest.warns(FutureWarning, match="AI_JUDGE is deprecated. Use LLM_JUDGE instead."):
73          source = AssessmentSource(source_type="AI_JUDGE", source_id="ai_1")
74  
75      assert source.source_type == "LLM_JUDGE"
76      assert source.source_id == "ai_1"