1  from unittest import mock
 2  
 3  import mlflow
 4  
 5  
 6  def test_doctor(capsys):
 7      mlflow.doctor()
 8      captured = capsys.readouterr()
 9      assert f"MLflow version: {mlflow.__version__}" in captured.out
10  
11  
12  def test_doctor_active_run(capsys):
13      with mlflow.start_run() as run:
14          mlflow.doctor()
15          captured = capsys.readouterr()
16          assert f"Active run ID: {run.info.run_id}" in captured.out
17  
18  
19  def test_doctor_databricks_runtime(capsys):
20      mock_version = "12.0"
21      with mock.patch(
22          "mlflow.utils.doctor.get_databricks_runtime_version", return_value=mock_version
23      ) as mock_runtime:
24          mlflow.doctor()
25          mock_runtime.assert_called_once()
26          captured = capsys.readouterr()
27          assert f"Databricks runtime version: {mock_version}" in captured.out