wehub-resource-sync
77 lines
1.8 KiB
Python
77 lines
1.8 KiB
Python
import threading
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.environment_variables import MLFLOW_TRACKING_USERNAME
|
|
from mlflow.genai.evaluation.context import NoneContext, _set_context, eval_context, get_context
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_context():
|
|
yield
|
|
_set_context(NoneContext())
|
|
|
|
|
|
def test_context_get_experiment_and_run_id():
|
|
exp_id = mlflow.set_experiment("Test").experiment_id
|
|
|
|
@eval_context
|
|
def _test():
|
|
assert exp_id == get_context().get_mlflow_experiment_id()
|
|
assert get_context().get_mlflow_run_id() is None
|
|
|
|
_test()
|
|
|
|
|
|
def test_context_get_run_id_active_run():
|
|
@eval_context
|
|
def _test():
|
|
with mlflow.start_run() as run:
|
|
assert run.info.run_id == get_context().get_mlflow_run_id()
|
|
|
|
_test()
|
|
|
|
|
|
def test_context_get_run_id_explicitly_set():
|
|
@eval_context
|
|
def _test():
|
|
context = get_context()
|
|
context.set_mlflow_run_id("test-run-id")
|
|
assert context.get_mlflow_run_id() == "test-run-id"
|
|
|
|
run_id = None
|
|
|
|
def _target():
|
|
nonlocal run_id
|
|
run_id = get_context().get_mlflow_run_id()
|
|
|
|
thread = threading.Thread(name="test-evaluate-context", target=_target)
|
|
thread.start()
|
|
assert run_id == "test-run-id"
|
|
|
|
_test()
|
|
|
|
|
|
def test_context_get_user_name(monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "test-user")
|
|
|
|
@eval_context
|
|
def _test():
|
|
assert get_context().get_user_name() == "test-user"
|
|
|
|
_test()
|
|
|
|
|
|
def test_context_get_user_name_no_user_set():
|
|
with mock.patch(
|
|
"mlflow.tracking.context.default_context.DefaultRunContext.tags", return_value={}
|
|
):
|
|
|
|
@eval_context
|
|
def _test():
|
|
assert get_context().get_user_name() == "unknown"
|
|
|
|
_test()
|