/ tests / data / test_code_dataset_source.py
test_code_dataset_source.py
 1  from mlflow.data.code_dataset_source import CodeDatasetSource
 2  
 3  
 4  def test_code_dataset_source_from_path():
 5      tags = {
 6          "mlflow_source_type": "NOTEBOOK",
 7          "mlflow_source_name": "some_random_notebook_path",
 8      }
 9      code_datasource = CodeDatasetSource(tags)
10      assert code_datasource.to_dict() == {
11          "tags": tags,
12      }
13  
14  
15  def test_code_datasource_type():
16      assert CodeDatasetSource._get_source_type() == "code"