1  import json
 2  
 3  import pytest
 4  
 5  from mlflow.genai.datasets.databricks_evaluation_dataset_source import (
 6      DatabricksEvaluationDatasetSource,
 7  )
 8  
 9  
10  def test_databricks_evaluation_dataset_source_init():
11      source = DatabricksEvaluationDatasetSource(
12          table_name="catalog.schema.table", dataset_id="12345"
13      )
14      assert source.table_name == "catalog.schema.table"
15      assert source.dataset_id == "12345"
16  
17  
18  def test_databricks_evaluation_dataset_source_get_source_type():
19      assert DatabricksEvaluationDatasetSource._get_source_type() == "databricks_evaluation_dataset"
20  
21  
22  def test_databricks_evaluation_dataset_source_to_dict():
23      source = DatabricksEvaluationDatasetSource(
24          table_name="catalog.schema.table", dataset_id="12345"
25      )
26      assert source.to_dict() == {
27          "table_name": "catalog.schema.table",
28          "dataset_id": "12345",
29      }
30  
31  
32  def test_databricks_evaluation_dataset_source_from_dict():
33      source_dict = {"table_name": "catalog.schema.table", "dataset_id": "12345"}
34      source = DatabricksEvaluationDatasetSource.from_dict(source_dict)
35      assert source.table_name == "catalog.schema.table"
36      assert source.dataset_id == "12345"
37  
38  
39  def test_databricks_evaluation_dataset_source_to_json():
40      source = DatabricksEvaluationDatasetSource(
41          table_name="catalog.schema.table", dataset_id="12345"
42      )
43      json_str = source.to_json()
44      parsed = json.loads(json_str)
45      assert parsed == {"table_name": "catalog.schema.table", "dataset_id": "12345"}
46  
47  
48  def test_databricks_evaluation_dataset_source_from_json():
49      json_str = json.dumps({"table_name": "catalog.schema.table", "dataset_id": "12345"})
50      source = DatabricksEvaluationDatasetSource.from_json(json_str)
51      assert source.table_name == "catalog.schema.table"
52      assert source.dataset_id == "12345"
53  
54  
55  def test_databricks_evaluation_dataset_source_load_not_implemented():
56      source = DatabricksEvaluationDatasetSource(
57          table_name="catalog.schema.table", dataset_id="12345"
58      )
59      with pytest.raises(
60          NotImplementedError,
61          match="Loading a Databricks Evaluation Dataset from source is not supported",
62      ):
63          source.load()
64  
65  
66  def test_databricks_evaluation_dataset_source_can_resolve():
67      # _can_resolve should return False for all inputs
68      assert DatabricksEvaluationDatasetSource._can_resolve({}) is False
69      assert DatabricksEvaluationDatasetSource._can_resolve({"table_name": "test"}) is False
70  
71  
72  def test_databricks_evaluation_dataset_source_resolve_not_implemented():
73      with pytest.raises(
74          NotImplementedError, match="Resolution from a source dictionary is not supported"
75      ):
76          DatabricksEvaluationDatasetSource._resolve({})