wehub-resource-sync
81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
from unittest import mock
|
|
|
|
import mlflow
|
|
from mlflow.genai.utils import display_utils
|
|
from mlflow.store.tracking.rest_store import RestStore
|
|
from mlflow.tracking.client import MlflowClient
|
|
from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_WORKSPACE_URL
|
|
|
|
|
|
def test_display_outputs_jupyter(monkeypatch):
|
|
mock_store = mock.MagicMock(spec=RestStore)
|
|
mock_store.get_run = MlflowClient().get_run
|
|
mock_store.get_host_creds = lambda: mock.MagicMock(host="https://mlflow.example.com/")
|
|
|
|
with (
|
|
mock.patch("IPython.display.display") as mock_display,
|
|
mock.patch.object(display_utils, "_get_store", return_value=mock_store),
|
|
mock.patch.object(display_utils, "_is_jupyter", return_value=True),
|
|
mlflow.start_run() as run,
|
|
):
|
|
display_utils.display_evaluation_output(run.info.run_id)
|
|
|
|
exp_id = run.info.experiment_id
|
|
expected_url = f"https://mlflow.example.com/#/experiments/{exp_id}/evaluation-runs?selectedRunUuid={run.info.run_id}"
|
|
html_content = mock_display.call_args[0][0].data
|
|
assert expected_url in html_content
|
|
|
|
|
|
def test_display_outputs_non_ipython(capsys):
|
|
mock_store = mock.MagicMock(spec=RestStore)
|
|
mock_store.get_run = mlflow.tracking.MlflowClient().get_run
|
|
mock_store.get_host_creds = lambda: mock.MagicMock(host="https://mlflow.example.com/")
|
|
|
|
with (
|
|
mock.patch.object(display_utils, "_get_store", return_value=mock_store),
|
|
mock.patch.object(display_utils, "_is_jupyter", return_value=False),
|
|
mlflow.start_run() as run,
|
|
):
|
|
display_utils.display_evaluation_output(run.info.run_id)
|
|
|
|
captured = capsys.readouterr().out
|
|
exp_id = run.info.experiment_id
|
|
expected_url = f"https://mlflow.example.com/#/experiments/{exp_id}/evaluation-runs?selectedRunUuid={run.info.run_id}"
|
|
assert expected_url in captured
|
|
|
|
|
|
def test_display_outputs_databricks(monkeypatch):
|
|
host = "https://workspace.databricks.com"
|
|
client = mlflow.tracking.MlflowClient()
|
|
|
|
mock_store = mock.MagicMock(spec=RestStore)
|
|
mock_store.get_run = client.get_run
|
|
mock_store.get_host_creds = lambda: mock.MagicMock(host=host)
|
|
|
|
with mlflow.start_run() as run:
|
|
client.set_tag(run.info.run_id, MLFLOW_DATABRICKS_WORKSPACE_URL, host)
|
|
|
|
with (
|
|
mock.patch("IPython.display.display") as mock_display,
|
|
mock.patch.object(display_utils, "_get_store", return_value=mock_store),
|
|
mock.patch.object(display_utils, "_is_jupyter", return_value=True),
|
|
mock.patch.object(display_utils, "is_databricks_uri", return_value=True),
|
|
):
|
|
display_utils.display_evaluation_output(run.info.run_id)
|
|
|
|
exp_id = run.info.experiment_id
|
|
expected_url = (
|
|
f"{host}/ml/experiments/{exp_id}/evaluation-runs?selectedRunUuid={run.info.run_id}"
|
|
)
|
|
html_content = mock_display.call_args[0][0].data
|
|
assert expected_url in html_content
|
|
|
|
|
|
def test_display_summary_with_local_store(capsys):
|
|
with mlflow.start_run() as run:
|
|
display_utils.display_evaluation_output(run.info.run_id)
|
|
|
|
captured = capsys.readouterr().out
|
|
assert run.info.run_id in captured
|
|
assert "Traces" in captured
|