mlflow--mlflow
wehub-resource-sync
43 行
1.4 KiB
Python
43 行
1.4 KiB
Python
from unittest import mock
|
|
|
|
from mlflow import set_experiment
|
|
from mlflow.demo.base import DEMO_EXPERIMENT_NAME
|
|
from mlflow.demo.generators.evaluation import EvaluationDemoGenerator
|
|
from mlflow.demo.generators.issues import IssuesDemoGenerator
|
|
from mlflow.demo.generators.traces import TracesDemoGenerator
|
|
|
|
|
|
def test_evaluation_fetch_demo_traces_passes_flush():
|
|
with mock.patch(
|
|
"mlflow.demo.generators.evaluation.mlflow.search_traces",
|
|
return_value=[],
|
|
) as mock_search:
|
|
EvaluationDemoGenerator()._fetch_demo_traces(experiment_id="0", version="v2", session=True)
|
|
|
|
mock_search.assert_called_once()
|
|
assert mock_search.call_args.kwargs.get("flush") is True
|
|
|
|
|
|
def test_issues_generate_passes_flush():
|
|
set_experiment(DEMO_EXPERIMENT_NAME)
|
|
with mock.patch(
|
|
"mlflow.demo.generators.issues.mlflow.search_traces",
|
|
return_value=[],
|
|
) as mock_search:
|
|
IssuesDemoGenerator().generate()
|
|
|
|
mock_search.assert_called_once()
|
|
assert mock_search.call_args.kwargs.get("flush") is True
|
|
|
|
|
|
def test_traces_data_exists_passes_flush():
|
|
set_experiment(DEMO_EXPERIMENT_NAME)
|
|
with mock.patch(
|
|
"mlflow.demo.generators.traces.mlflow.search_traces",
|
|
return_value=[],
|
|
) as mock_search:
|
|
TracesDemoGenerator()._data_exists()
|
|
|
|
mock_search.assert_called_once()
|
|
assert mock_search.call_args.kwargs.get("flush") is True
|