import json import urllib.parse import uuid from dataclasses import asdict from datetime import datetime, timezone from unittest import mock import pytest from flask import Response from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan import mlflow from mlflow.entities import ( Experiment, GatewayBudgetPolicy, Issue, IssueSeverity, IssueStatus, RunStatus, ScorerVersion, Span, Trace, TraceData, TraceInfo, TraceState, ViewType, ) from mlflow.entities._job import Job as JobEntity from mlflow.entities._job_status import JobStatus from mlflow.entities.gateway_budget_policy import ( BudgetAction, BudgetDuration, BudgetDurationUnit, BudgetTargetScope, BudgetUnit, ) from mlflow.entities.model_registry import ( ModelVersion, ModelVersionTag, PromptVersion, RegisteredModel, RegisteredModelTag, ) from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY, PROMPT_TEXT_TAG_KEY from mlflow.entities.presigned_download import PresignedDownloadUrlResponse from mlflow.entities.presigned_upload import CreatePresignedUploadResponse from mlflow.entities.trace_location import TraceLocation as EntityTraceLocation from mlflow.entities.trace_metrics import ( AggregationType, MetricAggregation, MetricDataPoint, MetricViewType, ) from mlflow.environment_variables import MLFLOW_ENABLE_WORKSPACES from mlflow.exceptions import ( MlflowException, MlflowNotImplementedException, MlflowTracingException, ) from mlflow.gateway.budget_tracker.in_memory import InMemoryBudgetTracker from mlflow.genai.review_queues import ReviewQueueType from mlflow.genai.review_queues.review_queues import ( ReviewItemType, ReviewQueue, ReviewQueueItem, ReviewStatus, ) from mlflow.genai.scorers.online.entities import OnlineScoringConfig from mlflow.protos.databricks_pb2 import ( INTERNAL_ERROR, INVALID_PARAMETER_VALUE, NOT_IMPLEMENTED, RESOURCE_DOES_NOT_EXIST, ErrorCode, ) from mlflow.protos.issues_pb2 import ( CreateIssue, SearchIssues, UpdateIssue, ) from mlflow.protos.model_registry_pb2 import ( CreateModelVersion, CreateRegisteredModel, DeleteModelVersion, DeleteModelVersionTag, DeleteRegisteredModel, DeleteRegisteredModelAlias, DeleteRegisteredModelTag, GetLatestVersions, GetModelVersion, GetModelVersionByAlias, GetModelVersionDownloadUri, GetRegisteredModel, RenameRegisteredModel, SearchModelVersions, SearchRegisteredModels, SetModelVersionTag, SetRegisteredModelAlias, SetRegisteredModelTag, TransitionModelVersionStage, UpdateModelVersion, UpdateRegisteredModel, ) from mlflow.protos.prompt_optimization_pb2 import ( OPTIMIZER_TYPE_GEPA, OPTIMIZER_TYPE_METAPROMPT, OPTIMIZER_TYPE_UNSPECIFIED, ) from mlflow.protos.review_queues_pb2 import ( CreateReviewQueue, GetOrCreateUserQueue, SetReviewQueueItemStatus, UpdateReviewQueue, ) from mlflow.protos.service_pb2 import ( BatchGetTraceInfos, BatchGetTraces, CalculateTraceFilterCorrelation, CreateExperiment, DeleteScorer, DeleteTraceTag, DeleteTraceTagV3, GatewayEndpoint, GetGatewayEndpoint, GetScorer, GetTrace, LinkPromptsToTrace, ListScorers, ListScorerVersions, QueryTraceMetrics, RegisterScorer, SearchExperiments, SearchLoggedModels, SearchRuns, SearchTraces, SearchTracesV3, SetTraceTag, SetTraceTagV3, TraceLocation, ) from mlflow.protos.webhooks_pb2 import ListWebhooks from mlflow.server import ( ARTIFACTS_DESTINATION_ENV_VAR, BACKEND_STORE_URI_ENV_VAR, SERVE_ARTIFACTS_ENV_VAR, app, ) from mlflow.server.handlers import ( ARTIFACT_STREAM_CHUNK_SIZE, STATIC_PREFIX_ENV_VAR, ModelRegistryStoreRegistryWrapper, TrackingStoreRegistryWrapper, _batch_get_trace_infos, _batch_get_traces, _calculate_trace_filter_correlation, _cancel_prompt_optimization_job, _convert_path_parameter_to_flask_format, _create_artifact_file_response, _create_dataset_handler, _create_experiment, _create_issue, _create_model_version, _create_presigned_upload_url, _create_prompt_optimization_job, _create_registered_model, _create_review_queue, _delete_artifact_mlflow_artifacts, _delete_dataset_handler, _delete_dataset_tag_handler, _delete_model_version, _delete_model_version_tag, _delete_registered_model, _delete_registered_model_alias, _delete_registered_model_tag, _delete_scorer, _delete_trace_tag, _delete_trace_tag_v3, _deprecated_search_traces_v2, _download_artifact, _get_ajax_path, _get_dataset_experiment_ids_handler, _get_dataset_handler, _get_dataset_records_handler, _get_gateway_endpoint, _get_issue, _get_latest_versions, _get_model_version, _get_model_version_by_alias, _get_model_version_download_uri, _get_or_create_user_queue, _get_presigned_download_url, _get_registered_model, _get_request_message, _get_rest_path, _get_scorer, _get_trace, _get_trace_artifact_repo, _get_workspace_scoped_repo_path_if_enabled, _link_prompts_to_trace, _list_artifacts_for_proxied_run_artifact_root, _list_scorer_versions, _list_scorers, _list_webhooks, _log_batch, _query_trace_metrics, _register_scorer, _rename_registered_model, _response_with_file_attachment_headers, _search_evaluation_datasets_handler, _search_experiments, _search_issues, _search_logged_models, _search_model_versions, _search_registered_models, _search_runs, _search_traces_v3, _send_artifact, _set_dataset_tags_handler, _set_model_version_tag, _set_registered_model_alias, _set_registered_model_tag, _set_review_queue_item_status, _set_trace_tag, _set_trace_tag_v3, _transition_stage, _update_issue, _update_model_version, _update_registered_model, _update_review_queue, _upload_artifact, _upsert_dataset_records_handler, _validate_source_run, catch_mlflow_exception, get_artifact_handler, get_endpoints, get_logged_model_artifact_handler, get_model_version_artifact_handler, get_trace_artifact_handler, get_ui_telemetry_handler, post_ui_telemetry_handler, upload_artifact_handler, ) from mlflow.store._unity_catalog.registry.rest_store import UcModelRegistryStore from mlflow.store.artifact.artifact_repo import ArtifactRepository from mlflow.store.artifact.azure_blob_artifact_repo import AzureBlobArtifactRepository from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository from mlflow.store.entities.paged_list import PagedList from mlflow.store.model_registry import ( SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD, SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT, ) from mlflow.store.model_registry.rest_store import RestStore as ModelRegistryRestStore from mlflow.store.tracking import MAX_RESULTS_QUERY_TRACE_METRICS from mlflow.store.tracking.databricks_rest_store import DatabricksTracingRestStore from mlflow.telemetry.schemas import Record, Status from mlflow.tracing.analysis import TraceFilterCorrelationResult from mlflow.tracing.constant import SpansLocation, TraceTagKey from mlflow.tracing.utils import build_otel_context from mlflow.utils.mlflow_tags import MLFLOW_ARTIFACT_LOCATION from mlflow.utils.proto_json_utils import message_to_json from mlflow.utils.validation import MAX_BATCH_LOG_REQUEST_SIZE from mlflow.utils.workspace_context import WorkspaceContext from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME @pytest.fixture def mock_get_request_message(): with mock.patch("mlflow.server.handlers._get_request_message") as m: yield m @pytest.fixture def mock_get_request_json(): with mock.patch("mlflow.server.handlers._get_request_json") as m: yield m @pytest.fixture def mock_tracking_store(): with mock.patch("mlflow.server.handlers._get_tracking_store") as m: mock_store = mock.MagicMock() m.return_value = mock_store yield mock_store @pytest.fixture def mock_model_registry_store(): with mock.patch("mlflow.server.handlers._get_model_registry_store") as m: mock_store = mock.MagicMock() mock_store.list_webhooks_by_event.return_value = PagedList([], None) m.return_value = mock_store yield mock_store @pytest.fixture def enable_serve_artifacts(monkeypatch): monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") @pytest.fixture def mock_evaluation_dataset(): from mlflow.protos.datasets_pb2 import Dataset as ProtoDataset dataset = mock.MagicMock() dataset.dataset_id = "d-1234567890abcdef1234567890abcdef" dataset.name = "test_dataset" dataset.digest = "abc123" dataset.created_time = 1234567890 dataset.last_update_time = 1234567890 dataset.created_by = "test_user" dataset.last_updated_by = "test_user" dataset.tags = {"env": "test", "version": "1.0"} dataset.experiment_ids = ["0", "1"] dataset._records = [] dataset.schema = json.dumps({ "inputs": {"question": "string"}, "expectations": {"accuracy": "float"}, }) dataset.profile = json.dumps({"record_count": 0}) proto_dataset = ProtoDataset() proto_dataset.dataset_id = dataset.dataset_id proto_dataset.name = dataset.name proto_dataset.digest = dataset.digest proto_dataset.created_time = dataset.created_time proto_dataset.last_update_time = dataset.last_update_time proto_dataset.created_by = dataset.created_by proto_dataset.last_updated_by = dataset.last_updated_by proto_dataset.schema = dataset.schema proto_dataset.profile = dataset.profile dataset.to_proto = mock.MagicMock(return_value=proto_dataset) return dataset @pytest.fixture def mock_telemetry_config_cache(): with mock.patch("mlflow.server.handlers._telemetry_config_cache", {}) as m: yield m @pytest.fixture def bypass_telemetry_env_check(monkeypatch): monkeypatch.setattr(mlflow.telemetry.utils, "_IS_MLFLOW_TESTING_TELEMETRY", False) monkeypatch.setattr(mlflow.telemetry.utils, "_IS_IN_CI_ENV_OR_TESTING", False) monkeypatch.setattr(mlflow.telemetry.utils, "_IS_MLFLOW_DEV_VERSION", False) @pytest.fixture def mock_job_store(): with mock.patch("mlflow.server.handlers._get_job_store") as m: mock_store = mock.MagicMock() m.return_value = mock_store yield mock_store def _create_mock_job( job_id="job-123", job_name="optimize_prompts", status_name="PENDING", params=None, result=None, creation_time=1234567890000, status_details=None, ): from mlflow.entities._job import Job from mlflow.entities._job_status import JobStatus if params is None: params = { "experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1", "run_id": "run-456", } return Job( job_id=job_id, creation_time=creation_time, job_name=job_name, params=json.dumps(params), timeout=None, status=JobStatus.from_str(status_name), result=json.dumps(result) if result and status_name == "SUCCEEDED" else result, retry_count=0, last_update_time=creation_time, status_details=status_details, ) def _create_mock_run(run_id="run-456", params=None, metrics=None): mock_run = mock.MagicMock() mock_run.info.run_id = run_id mock_run.data.params = params or {} mock_run.data.metrics = metrics or {} return mock_run def test_health(): with app.test_client() as c: response = c.get("/health") assert response.status_code == 200 assert response.get_data().decode() == "OK" def test_version(): with app.test_client() as c: response = c.get("/version") assert response.status_code == 200 assert response.get_data().decode() == mlflow.__version__ def test_server_info(): with app.test_client() as c: response = c.get("/api/3.0/mlflow/server-info") assert response.status_code == 200 data = response.get_json() assert data["store_type"] == "SqlStore" assert data["workspaces_enabled"] is False assert data["trace_archival_enabled"] is False def test_server_info_trace_archival_enabled(monkeypatch): monkeypatch.setattr( "mlflow.server.handlers.get_trace_archival_server_config", mock.Mock(return_value=mock.Mock(enabled=True)), ) monkeypatch.setattr("mlflow.server.handlers._store_supports_trace_archival", lambda store: True) with app.test_client() as c: response = c.get("/api/3.0/mlflow/server-info") assert response.status_code == 200 data = response.get_json() assert data["trace_archival_enabled"] is True def test_server_info_handles_invalid_trace_archival_config(monkeypatch): monkeypatch.setattr( "mlflow.server.handlers.get_trace_archival_server_config", mock.Mock( side_effect=MlflowException.invalid_parameter_value("invalid trace archival config") ), ) with app.test_client() as c: response = c.get("/api/3.0/mlflow/server-info") assert response.status_code == 200 data = response.get_json() assert data["trace_archival_enabled"] is False def test_server_info_handles_unexpected_trace_archival_config_error(monkeypatch): monkeypatch.setattr( "mlflow.server.handlers.get_trace_archival_server_config", mock.Mock(side_effect=RuntimeError("unexpected trace archival config error")), ) with app.test_client() as c: response = c.get("/api/3.0/mlflow/server-info") assert response.status_code == 200 data = response.get_json() assert data["trace_archival_enabled"] is False def test_get_endpoints(): endpoints = get_endpoints() create_experiment_endpoint = [e for e in endpoints if e[1] == _create_experiment] assert len(create_experiment_endpoint) == 2 def test_convert_path_parameter_to_flask_format(): converted = _convert_path_parameter_to_flask_format("/mlflow/trace") assert "/mlflow/trace" == converted converted = _convert_path_parameter_to_flask_format("/mlflow/trace/{request_id}") assert "/mlflow/trace/" == converted converted = _convert_path_parameter_to_flask_format("/mlflow/{foo}/{bar}/{baz}") assert "/mlflow///" == converted def test_all_model_registry_endpoints_available(): endpoints = {handler: method for (path, handler, method) in get_endpoints()} # Test that each of the handler is enabled as an endpoint with appropriate method. expected_endpoints = { "POST": [ _create_registered_model, _create_model_version, _rename_registered_model, _transition_stage, ], "PATCH": [_update_registered_model, _update_model_version], "DELETE": [_delete_registered_model, _delete_registered_model], "GET": [ _search_model_versions, _get_latest_versions, _get_registered_model, _get_model_version, _get_model_version_download_uri, ], } # TODO: efficient mechanism to test endpoint path for method, handlers in expected_endpoints.items(): for handler in handlers: assert handler in endpoints assert endpoints[handler] == [method] def test_can_parse_json(): request = mock.MagicMock() request.method = "POST" request.content_type = "application/json" request.get_json = mock.MagicMock() request.get_json.return_value = {"name": "hello"} msg = _get_request_message(CreateExperiment(), flask_request=request) assert msg.name == "hello" def test_can_parse_post_json_with_unknown_fields(): request = mock.MagicMock() request.method = "POST" request.content_type = "application/json" request.get_json = mock.MagicMock() request.get_json.return_value = {"name": "hello", "WHAT IS THIS FIELD EVEN": "DOING"} msg = _get_request_message(CreateExperiment(), flask_request=request) assert msg.name == "hello" def test_can_parse_post_json_with_content_type_params(): request = mock.MagicMock() request.method = "POST" request.content_type = "application/json; charset=utf-8" request.get_json = mock.MagicMock() request.get_json.return_value = {"name": "hello"} msg = _get_request_message(CreateExperiment(), flask_request=request) assert msg.name == "hello" def test_can_parse_get_json_with_unknown_fields(): request = mock.MagicMock() request.method = "GET" request.args = {"name": "hello", "superDuperUnknown": "field"} msg = _get_request_message(CreateExperiment(), flask_request=request) assert msg.name == "hello" # Previous versions of the client sent a doubly string encoded JSON blob, # so this test ensures continued compliance with such clients. def test_can_parse_json_string(): request = mock.MagicMock() request.method = "POST" request.content_type = "application/json" request.get_json = mock.MagicMock() request.get_json.return_value = '{"name": "hello2"}' msg = _get_request_message(CreateExperiment(), flask_request=request) assert msg.name == "hello2" def test_can_block_post_request_with_invalid_content_type(): request = mock.MagicMock() request.method = "POST" request.content_type = "text/plain" request.get_json = mock.MagicMock() request.get_json.return_value = {"name": "hello"} with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"): _get_request_message(CreateExperiment(), flask_request=request) def test_can_block_post_request_with_missing_content_type(): request = mock.MagicMock() request.method = "POST" request.content_type = None request.get_json = mock.MagicMock() request.get_json.return_value = {"name": "hello"} with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"): _get_request_message(CreateExperiment(), flask_request=request) def test_search_runs_default_view_type(mock_get_request_message, mock_tracking_store): """ Search Runs default view type is filled in as ViewType.ACTIVE_ONLY """ mock_get_request_message.return_value = SearchRuns(experiment_ids=["0"]) mock_tracking_store.search_runs.return_value = PagedList([], None) _search_runs() _, kwargs = mock_tracking_store.search_runs.call_args assert kwargs["run_view_type"] == ViewType.ACTIVE_ONLY def test_search_runs_empty_page_token(mock_get_request_message, mock_tracking_store): """ Test that empty page_token from protobuf is converted to None before calling store """ # Create proto without setting page_token search_runs_proto = SearchRuns() search_runs_proto.experiment_ids.append("0") search_runs_proto.max_results = 10 # Verify protobuf returns empty string for unset field assert search_runs_proto.page_token == "" mock_get_request_message.return_value = search_runs_proto mock_tracking_store.search_runs.return_value = PagedList([], None) _search_runs() # Verify store was called with None, not empty string mock_tracking_store.search_runs.assert_called_once() call_kwargs = mock_tracking_store.search_runs.call_args.kwargs assert call_kwargs["page_token"] is None # page_token should be None, not "" def test_log_batch_api_req(mock_get_request_json): mock_get_request_json.return_value = "a" * (MAX_BATCH_LOG_REQUEST_SIZE + 1) response = _log_batch() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert ( f"Batched logging API requests must be at most {MAX_BATCH_LOG_REQUEST_SIZE} bytes" in json_response["message"] ) def test_catch_mlflow_exception(): @catch_mlflow_exception def test_handler(): raise MlflowException("test error", error_code=INTERNAL_ERROR) response = test_handler() json_response = json.loads(response.get_data()) assert response.status_code == 500 assert json_response["error_code"] == ErrorCode.Name(INTERNAL_ERROR) assert json_response["message"] == "test error" def test_mlflow_server_with_installed_plugin(tmp_path, monkeypatch): pytest.skip("FileStore is no longer supported.") from mlflow_test_plugin.file_store import PluginFileStore monkeypatch.setenv(BACKEND_STORE_URI_ENV_VAR, f"file-plugin:{tmp_path}") monkeypatch.setattr(mlflow.server.handlers, "_tracking_store", None) plugin_file_store = mlflow.server.handlers._get_tracking_store() assert isinstance(plugin_file_store, PluginFileStore) assert plugin_file_store.is_plugin def jsonify(obj): def _jsonify(obj): return json.loads(message_to_json(obj.to_proto())) if isinstance(obj, list): return [_jsonify(o) for o in obj] else: return _jsonify(obj) # Tests for Model Registry handlers def test_create_registered_model(mock_get_request_message, mock_model_registry_store): tags = [ RegisteredModelTag(key="key", value="value"), RegisteredModelTag(key="anotherKey", value="some other value"), ] mock_get_request_message.return_value = CreateRegisteredModel( name="model_1", tags=[tag.to_proto() for tag in tags] ) rm = RegisteredModel("model_1", tags=tags) mock_model_registry_store.create_registered_model.return_value = rm resp = _create_registered_model() _, args = mock_model_registry_store.create_registered_model.call_args assert args["name"] == "model_1" assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm)} def test_get_registered_model(mock_get_request_message, mock_model_registry_store): name = "model1" mock_get_request_message.return_value = GetRegisteredModel(name=name) rmd = RegisteredModel( name=name, creation_timestamp=111, last_updated_timestamp=222, description="Test model", latest_versions=[], ) mock_model_registry_store.get_registered_model.return_value = rmd resp = _get_registered_model() _, args = mock_model_registry_store.get_registered_model.call_args assert args == {"name": name} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rmd)} def test_update_registered_model(mock_get_request_message, mock_model_registry_store): name = "model_1" description = "Test model" mock_get_request_message.return_value = UpdateRegisteredModel( name=name, description=description ) rm2 = RegisteredModel(name, description=description) mock_model_registry_store.update_registered_model.return_value = rm2 resp = _update_registered_model() _, args = mock_model_registry_store.update_registered_model.call_args assert args == {"name": name, "description": "Test model"} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm2)} def test_rename_registered_model(mock_get_request_message, mock_model_registry_store): name = "model_1" new_name = "model_2" mock_get_request_message.return_value = RenameRegisteredModel(name=name, new_name=new_name) rm2 = RegisteredModel(new_name) mock_model_registry_store.rename_registered_model.return_value = rm2 resp = _rename_registered_model() _, args = mock_model_registry_store.rename_registered_model.call_args assert args == {"name": name, "new_name": new_name} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm2)} def test_delete_registered_model(mock_get_request_message, mock_model_registry_store): name = "model_1" mock_get_request_message.return_value = DeleteRegisteredModel(name=name) _delete_registered_model() _, args = mock_model_registry_store.delete_registered_model.call_args assert args == {"name": name} def test_search_registered_models(mock_get_request_message, mock_model_registry_store): rmds = [ RegisteredModel( name="model_1", creation_timestamp=111, last_updated_timestamp=222, description="Test model", latest_versions=[], ), RegisteredModel( name="model_2", creation_timestamp=111, last_updated_timestamp=333, description="Another model", latest_versions=[], ), ] mock_get_request_message.return_value = SearchRegisteredModels() mock_model_registry_store.search_registered_models.return_value = PagedList(rmds, None) resp = _search_registered_models() _, args = mock_model_registry_store.search_registered_models.call_args assert args == { "filter_string": "", "max_results": SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT, "order_by": [], "page_token": None, } assert json.loads(resp.get_data()) == {"registered_models": jsonify(rmds)} mock_get_request_message.return_value = SearchRegisteredModels(filter="hello") mock_model_registry_store.search_registered_models.return_value = PagedList(rmds[:1], "tok") resp = _search_registered_models() _, args = mock_model_registry_store.search_registered_models.call_args assert args == { "filter_string": "hello", "max_results": SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT, "order_by": [], "page_token": None, } assert json.loads(resp.get_data()) == { "registered_models": jsonify(rmds[:1]), "next_page_token": "tok", } mock_get_request_message.return_value = SearchRegisteredModels(filter="hi", max_results=5) mock_model_registry_store.search_registered_models.return_value = PagedList([rmds[0]], "tik") resp = _search_registered_models() _, args = mock_model_registry_store.search_registered_models.call_args assert args == {"filter_string": "hi", "max_results": 5, "order_by": [], "page_token": None} assert json.loads(resp.get_data()) == { "registered_models": jsonify([rmds[0]]), "next_page_token": "tik", } mock_get_request_message.return_value = SearchRegisteredModels( filter="hey", max_results=500, order_by=["a", "B desc"], page_token="prev" ) mock_model_registry_store.search_registered_models.return_value = PagedList(rmds, "DONE") resp = _search_registered_models() _, args = mock_model_registry_store.search_registered_models.call_args assert args == { "filter_string": "hey", "max_results": 500, "order_by": ["a", "B desc"], "page_token": "prev", } assert json.loads(resp.get_data()) == { "registered_models": jsonify(rmds), "next_page_token": "DONE", } def test_get_latest_versions(mock_get_request_message, mock_model_registry_store): name = "model1" mock_get_request_message.return_value = GetLatestVersions(name=name) mvds = [ ModelVersion( name=name, version="5", creation_timestamp=1, last_updated_timestamp=12, description="v 5", user_id="u1", current_stage="Production", source="A/B", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ModelVersion( name=name, version="1", creation_timestamp=1, last_updated_timestamp=1200, description="v 1", user_id="u1", current_stage="Archived", source="A/B2", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ModelVersion( name=name, version="12", creation_timestamp=100, last_updated_timestamp=None, description="v 12", user_id="u2", current_stage="Staging", source="A/B3", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ] mock_model_registry_store.get_latest_versions.return_value = mvds resp = _get_latest_versions() _, args = mock_model_registry_store.get_latest_versions.call_args assert args == {"name": name, "stages": []} assert json.loads(resp.get_data()) == {"model_versions": jsonify(mvds)} for stages in [[], ["None"], ["Staging"], ["Staging", "Production"]]: mock_get_request_message.return_value = GetLatestVersions(name=name, stages=stages) _get_latest_versions() _, args = mock_model_registry_store.get_latest_versions.call_args assert args == {"name": name, "stages": stages} def test_create_model_version(mock_get_request_message, mock_model_registry_store): run_id = uuid.uuid4().hex tags = [ ModelVersionTag(key="key", value="value"), ModelVersionTag(key="anotherKey", value="some other value"), ] run_link = "localhost:5000/path/to/run" mock_get_request_message.return_value = CreateModelVersion( name="model_1", source=f"runs:/{run_id}", run_id=run_id, run_link=run_link, tags=[tag.to_proto() for tag in tags], ) mv = ModelVersion( name="model_1", version="12", creation_timestamp=123, tags=tags, run_link=run_link ) mock_model_registry_store.create_model_version.return_value = mv resp = _create_model_version() _, args = mock_model_registry_store.create_model_version.call_args assert args["name"] == "model_1" assert args["source"] == f"runs:/{run_id}" assert args["run_id"] == run_id assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert args["run_link"] == run_link assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)} @pytest.mark.parametrize( "source", [ "file:///etc/passwd", "file:///", "/etc/passwd", "file:///proc/self/environ", "file://remote-host/etc/passwd", "file://remote-host/", ], ) def test_create_model_version_rejects_local_source_for_prompts( mock_get_request_message, mock_model_registry_store, source ): mock_get_request_message.return_value = CreateModelVersion( name="model_1", source=source, tags=[ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true").to_proto()], ) resp = _create_model_version() assert resp.status_code == 400 assert "Invalid prompt source" in resp.get_json()["message"] @pytest.mark.parametrize( "source", [ "https://example.com/../../etc/passwd", "http://example.com/path/..%2f..%2fsecret", ], ) def test_create_model_version_rejects_traversal_source_for_prompts( mock_get_request_message, mock_model_registry_store, source ): mock_get_request_message.return_value = CreateModelVersion( name="model_1", source=source, tags=[ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true").to_proto()], ) resp = _create_model_version() assert resp.status_code == 400 assert "Invalid model version source" in resp.get_json()["message"] def test_set_registered_model_tag(mock_get_request_message, mock_model_registry_store): name = "model1" tag = RegisteredModelTag(key="some weird key", value="some value") mock_get_request_message.return_value = SetRegisteredModelTag( name=name, key=tag.key, value=tag.value ) _set_registered_model_tag() _, args = mock_model_registry_store.set_registered_model_tag.call_args assert args == {"name": name, "tag": tag} def test_delete_registered_model_tag(mock_get_request_message, mock_model_registry_store): name = "model1" key = "some weird key" mock_get_request_message.return_value = DeleteRegisteredModelTag(name=name, key=key) _delete_registered_model_tag() _, args = mock_model_registry_store.delete_registered_model_tag.call_args assert args == {"name": name, "key": key} def test_get_model_version_details(mock_get_request_message, mock_model_registry_store): mock_get_request_message.return_value = GetModelVersion(name="model1", version="32") mvd = ModelVersion( name="model1", version="5", creation_timestamp=1, last_updated_timestamp=12, description="v 5", user_id="u1", current_stage="Production", source="A/B", run_id=uuid.uuid4().hex, status="READY", status_message=None, ) mock_model_registry_store.get_model_version.return_value = mvd resp = _get_model_version() _, args = mock_model_registry_store.get_model_version.call_args assert args == {"name": "model1", "version": "32"} assert json.loads(resp.get_data()) == {"model_version": jsonify(mvd)} def test_update_model_version(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" description = "Great model!" mock_get_request_message.return_value = UpdateModelVersion( name=name, version=version, description=description ) mv = ModelVersion(name=name, version=version, creation_timestamp=123, description=description) mock_model_registry_store.update_model_version.return_value = mv _update_model_version() _, args = mock_model_registry_store.update_model_version.call_args assert args == {"name": name, "version": version, "description": description} def test_transition_model_version_stage(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" stage = "Production" mock_get_request_message.return_value = TransitionModelVersionStage( name=name, version=version, stage=stage ) mv = ModelVersion(name=name, version=version, creation_timestamp=123, current_stage=stage) mock_model_registry_store.transition_model_version_stage.return_value = mv _transition_stage() _, args = mock_model_registry_store.transition_model_version_stage.call_args assert args == { "name": name, "version": version, "stage": stage, "archive_existing_versions": False, } def test_delete_model_version(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" mock_get_request_message.return_value = DeleteModelVersion(name=name, version=version) _delete_model_version() _, args = mock_model_registry_store.delete_model_version.call_args assert args == {"name": name, "version": version} def test_get_model_version_download_uri(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" mock_get_request_message.return_value = GetModelVersionDownloadUri(name=name, version=version) mock_model_registry_store.get_model_version_download_uri.return_value = "some/download/path" resp = _get_model_version_download_uri() _, args = mock_model_registry_store.get_model_version_download_uri.call_args assert args == {"name": name, "version": version} assert json.loads(resp.get_data()) == {"artifact_uri": "some/download/path"} def test_search_model_versions(mock_get_request_message, mock_model_registry_store): mvds = [ ModelVersion( name="model_1", version="5", creation_timestamp=100, last_updated_timestamp=3200, description="v 5", user_id="u1", current_stage="Production", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ModelVersion( name="model_1", version="12", creation_timestamp=110, last_updated_timestamp=2000, description="v 12", user_id="u2", current_stage="Production", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ModelVersion( name="ads_model", version="8", creation_timestamp=200, last_updated_timestamp=1000, description="v 8", user_id="u1", current_stage="Staging", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ModelVersion( name="fraud_detection_model", version="345", creation_timestamp=1000, last_updated_timestamp=999, description="newest version", user_id="u12", current_stage="None", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None, ), ] mock_get_request_message.return_value = SearchModelVersions(filter="source_path = 'A/B/CD'") mock_model_registry_store.search_model_versions.return_value = PagedList(mvds, None) resp = _search_model_versions() mock_model_registry_store.search_model_versions.assert_called_with( filter_string="source_path = 'A/B/CD'", max_results=SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD, order_by=[], page_token=None, ) assert json.loads(resp.get_data()) == {"model_versions": jsonify(mvds)} mock_get_request_message.return_value = SearchModelVersions(filter="name='model_1'") mock_model_registry_store.search_model_versions.return_value = PagedList(mvds[:1], "tok") resp = _search_model_versions() mock_model_registry_store.search_model_versions.assert_called_with( filter_string="name='model_1'", max_results=SEARCH_MODEL_VERSION_MAX_RESULTS_THRESHOLD, order_by=[], page_token=None, ) assert json.loads(resp.get_data()) == { "model_versions": jsonify(mvds[:1]), "next_page_token": "tok", } mock_get_request_message.return_value = SearchModelVersions(filter="version<=12", max_results=2) mock_model_registry_store.search_model_versions.return_value = PagedList( [mvds[0], mvds[2]], "next" ) resp = _search_model_versions() mock_model_registry_store.search_model_versions.assert_called_with( filter_string="version<=12", max_results=2, order_by=[], page_token=None ) assert json.loads(resp.get_data()) == { "model_versions": jsonify([mvds[0], mvds[2]]), "next_page_token": "next", } mock_get_request_message.return_value = SearchModelVersions( filter="version<=12", max_results=2, order_by=["version DESC"], page_token="prev" ) mock_model_registry_store.search_model_versions.return_value = PagedList(mvds[1:3], "next") resp = _search_model_versions() mock_model_registry_store.search_model_versions.assert_called_with( filter_string="version<=12", max_results=2, order_by=["version DESC"], page_token="prev" ) assert json.loads(resp.get_data()) == { "model_versions": jsonify(mvds[1:3]), "next_page_token": "next", } def test_set_model_version_tag(mock_get_request_message, mock_model_registry_store): name = "model1" version = "1" tag = ModelVersionTag(key="some weird key", value="some value") mock_get_request_message.return_value = SetModelVersionTag( name=name, version=version, key=tag.key, value=tag.value ) _set_model_version_tag() _, args = mock_model_registry_store.set_model_version_tag.call_args assert args == {"name": name, "version": version, "tag": tag} def test_delete_model_version_tag(mock_get_request_message, mock_model_registry_store): name = "model1" version = "1" key = "some weird key" mock_get_request_message.return_value = DeleteModelVersionTag( name=name, version=version, key=key ) _delete_model_version_tag() _, args = mock_model_registry_store.delete_model_version_tag.call_args assert args == {"name": name, "version": version, "key": key} def test_set_registered_model_alias(mock_get_request_message, mock_model_registry_store): name = "model1" alias = "test_alias" version = "1" mock_get_request_message.return_value = SetRegisteredModelAlias( name=name, alias=alias, version=version ) _set_registered_model_alias() _, args = mock_model_registry_store.set_registered_model_alias.call_args assert args == {"name": name, "alias": alias, "version": version} def test_delete_registered_model_alias(mock_get_request_message, mock_model_registry_store): name = "model1" alias = "test_alias" mock_get_request_message.return_value = DeleteRegisteredModelAlias(name=name, alias=alias) _delete_registered_model_alias() _, args = mock_model_registry_store.delete_registered_model_alias.call_args assert args == {"name": name, "alias": alias} def test_get_model_version_by_alias(mock_get_request_message, mock_model_registry_store): name = "model1" alias = "test_alias" mock_get_request_message.return_value = GetModelVersionByAlias(name=name, alias=alias) mvd = ModelVersion( name="model1", version="5", creation_timestamp=1, last_updated_timestamp=12, description="v 5", user_id="u1", current_stage="Production", source="A/B", run_id=uuid.uuid4().hex, status="READY", status_message=None, aliases=["test_alias"], ) mock_model_registry_store.get_model_version_by_alias.return_value = mvd resp = _get_model_version_by_alias() _, args = mock_model_registry_store.get_model_version_by_alias.call_args assert args == {"name": name, "alias": alias} assert json.loads(resp.get_data()) == {"model_version": jsonify(mvd)} @pytest.mark.parametrize( "path", [ "/path", "path/../to/file", "/etc/passwd", "/etc/passwd%00.jpg", "/file://etc/passwd", "%2E%2E%2F%2E%2E%2Fpath", ], ) def test_delete_artifact_mlflow_artifacts_throws_for_malicious_path(enable_serve_artifacts, path): response = _delete_artifact_mlflow_artifacts(path) assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert json_response["message"] == "Invalid path" def test_get_presigned_download_url_success(enable_serve_artifacts): from mlflow.store.artifact.artifact_repo import MultipartDownloadMixin class MockMultipartDownloadRepo(MultipartDownloadMixin): def get_download_presigned_url(self, artifact_path, expiration=300): return PresignedDownloadUrlResponse( url="https://storage.example.com/presigned?token=abc", headers={"x-custom-header": "value"}, file_size=1024, ) artifact_path = "run_id/artifacts/model.pkl" with ( app.test_request_context(method="GET"), mock.patch( "mlflow.server.handlers._get_artifact_repo_mlflow_artifacts", return_value=MockMultipartDownloadRepo(), ), ): response = _get_presigned_download_url(artifact_path) assert response.status_code == 200 data = json.loads(response.get_data()) assert data["url"] == "https://storage.example.com/presigned?token=abc" assert data["headers"] == {"x-custom-header": "value"} assert data["file_size"] == 1024 @pytest.mark.parametrize( "path", [ "/path", "path/../to/file", "/etc/passwd", "/etc/passwd%00.jpg", "/file://etc/passwd", "%2E%2E%2F%2E%2E%2Fpath", ], ) def test_get_presigned_download_url_throws_for_malicious_path(enable_serve_artifacts, path): response = _get_presigned_download_url(path) assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert json_response["message"] == "Invalid path" def test_get_presigned_download_url_unsupported_repo(enable_serve_artifacts, tmp_path): with ( app.test_request_context(method="GET"), mock.patch( "mlflow.server.handlers._get_artifact_repo_mlflow_artifacts", return_value=LocalArtifactRepository(str(tmp_path)), ), ): response = _get_presigned_download_url("some/artifact/path") assert response.status_code == 501 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(NOT_IMPLEMENTED) assert "multipart" in json_response["message"].lower() # --- Presigned upload URL handler tests --- def test_create_presigned_upload_url_success(): from mlflow.store.artifact.artifact_repo import PresignedUploadMixin class MockPresignedUploadRepo(PresignedUploadMixin): def create_presigned_upload_url(self, artifact_path, expiration=900): return CreatePresignedUploadResponse( presigned_url="https://s3.amazonaws.com/bucket/artifacts/model.pkl?X-Amz-Signature=abc", headers={"Content-Type": "application/octet-stream"}, ) mock_run = mock.MagicMock() mock_run.info.artifact_uri = "s3://bucket/0/abc123/artifacts" from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = "model.pkl" with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, mock.patch( "mlflow.server.handlers._get_artifact_repo", return_value=MockPresignedUploadRepo(), ), ): mock_store.return_value.get_run.return_value = mock_run response = _create_presigned_upload_url() assert response.status_code == 200 data = json.loads(response.get_data()) assert "presigned_url" in data assert "X-Amz-Signature" in data["presigned_url"] assert data["headers"] == {"Content-Type": "application/octet-stream"} def test_create_presigned_upload_url_unsupported_repo(): mock_run = mock.MagicMock() mock_run.info.artifact_uri = "file:///tmp/artifacts" from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = "model.pkl" with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, mock.patch( "mlflow.server.handlers._get_artifact_repo", return_value=LocalArtifactRepository("/tmp/artifacts"), ), ): mock_store.return_value.get_run.return_value = mock_run response = _create_presigned_upload_url() assert response.status_code == 501 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(NOT_IMPLEMENTED) assert "presigned upload" in json_response["message"].lower() @pytest.mark.parametrize( "artifact_uri", [ "mlflow-artifacts:/0/abc123/artifacts", "http://mlflow-server:5000/api/2.0/mlflow-artifacts/artifacts", "https://mlflow-server/api/2.0/mlflow-artifacts/artifacts", ], ) def test_create_presigned_upload_url_rejects_proxy_artifact_uri(artifact_uri): mock_run = mock.MagicMock() mock_run.info.artifact_uri = artifact_uri from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = "model.pkl" with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, ): mock_store.return_value.get_run.return_value = mock_run response = _create_presigned_upload_url() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert "proxied" in json_response["message"].lower() def test_create_presigned_upload_url_invalid_run_id(): from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "nonexistent_run" request_proto.path = "model.pkl" with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, ): mock_store.return_value.get_run.side_effect = MlflowException( "Run 'nonexistent_run' not found", error_code=RESOURCE_DOES_NOT_EXIST, ) response = _create_presigned_upload_url() assert response.status_code == 404 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) @pytest.mark.parametrize( "path", [ "../../../etc/passwd", "path/../to/file", "/etc/passwd", "/etc/passwd%00.jpg", "%2E%2E%2F%2E%2E%2Fpath", ], ) def test_create_presigned_upload_url_rejects_path_traversal(path): from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = path with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), ): response = _create_presigned_upload_url() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) def test_create_presigned_upload_url_with_custom_expiration(): from mlflow.store.artifact.artifact_repo import PresignedUploadMixin captured_expiration = {} class MockPresignedUploadRepo(PresignedUploadMixin): def create_presigned_upload_url(self, artifact_path, expiration=900): captured_expiration["value"] = expiration return CreatePresignedUploadResponse( presigned_url="https://example.com/presigned", headers={}, ) mock_run = mock.MagicMock() mock_run.info.artifact_uri = "s3://bucket/0/abc123/artifacts" from mlflow.protos.service_pb2 import CreatePresignedUploadUrl request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = "model.pkl" request_proto.expiration = 60 with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, mock.patch( "mlflow.server.handlers._get_artifact_repo", return_value=MockPresignedUploadRepo(), ), ): mock_store.return_value.get_run.return_value = mock_run response = _create_presigned_upload_url() assert response.status_code == 200 assert captured_expiration["value"] == 60 def test_create_presigned_upload_url_default_expiration(): from mlflow.store.artifact.artifact_repo import PresignedUploadMixin captured_expiration = {} class MockPresignedUploadRepo(PresignedUploadMixin): def create_presigned_upload_url(self, artifact_path, expiration=900): captured_expiration["value"] = expiration return CreatePresignedUploadResponse( presigned_url="https://example.com/presigned", headers={}, ) mock_run = mock.MagicMock() mock_run.info.artifact_uri = "s3://bucket/0/abc123/artifacts" from mlflow.protos.service_pb2 import CreatePresignedUploadUrl # Don't set expiration - should default to 900 request_proto = CreatePresignedUploadUrl() request_proto.run_id = "abc123" request_proto.path = "model.pkl" with ( app.test_request_context(method="POST", content_type="application/json"), mock.patch( "mlflow.server.handlers._get_request_message", return_value=request_proto, ), mock.patch( "mlflow.server.handlers._get_tracking_store", ) as mock_store, mock.patch( "mlflow.server.handlers._get_artifact_repo", return_value=MockPresignedUploadRepo(), ), ): mock_store.return_value.get_run.return_value = mock_run response = _create_presigned_upload_url() assert response.status_code == 200 assert captured_expiration["value"] == 900 def test_create_presigned_upload_url_blocked_in_artifacts_only_mode(monkeypatch): from mlflow.server import ARTIFACTS_ONLY_ENV_VAR monkeypatch.setenv(ARTIFACTS_ONLY_ENV_VAR, "true") with app.test_request_context(method="POST", content_type="application/json"): response = _create_presigned_upload_url() assert response.status_code == 503 assert "artifacts-only" in response.get_data(as_text=True).lower() @pytest.mark.parametrize( "uri", [ "http://host#/abc/etc/", "http://host/;..%2F..%2Fetc", ], ) def test_local_file_read_write_by_pass_vulnerability(uri): request = mock.MagicMock() request.method = "POST" request.content_type = "application/json; charset=utf-8" request.get_json = mock.MagicMock() request.get_json.return_value = { "name": "hello", "artifact_location": uri, } msg = _get_request_message(CreateExperiment(), flask_request=request) with mock.patch("mlflow.server.handlers._get_request_message", return_value=msg): response = _create_experiment() json_response = json.loads(response.get_data()) assert ( json_response["message"] == "'artifact_location' URL can't include fragments or params." ) # Test if source is a local filesystem path, `_validate_source` validates that the run # artifact_uri is also a local filesystem path. run_id = uuid.uuid4().hex with mock.patch("mlflow.server.handlers._get_tracking_store") as mock_get_tracking_store: mock_get_tracking_store().get_run( run_id ).info.artifact_uri = f"http://host/{run_id}/artifacts/abc" with pytest.raises( MlflowException, match=( "the run_id request parameter has to be specified and the local " "path has to be contained within the artifact directory of the " "run specified by the run_id" ), ): _validate_source_run("/local/path/xyz", run_id) @pytest.mark.parametrize( ("location", "expected_class", "expected_uri"), [ ("file:///0/traces/123", LocalArtifactRepository, "file:///0/traces/123"), ("s3://bucket/0/traces/123", S3ArtifactRepository, "s3://bucket/0/traces/123"), ( "wasbs://container@account.blob.core.windows.net/bucket/1/traces/123", AzureBlobArtifactRepository, "wasbs://container@account.blob.core.windows.net/bucket/1/traces/123", ), # Proxy URI must be resolved to the actual storage URI ( "https://127.0.0.1/api/2.0/mlflow-artifacts/artifacts/2/traces/123", S3ArtifactRepository, "s3://bucket/2/traces/123", ), ("mlflow-artifacts:/1/traces/123", S3ArtifactRepository, "s3://bucket/1/traces/123"), ], ) def test_get_trace_artifact_repo(location, expected_class, expected_uri, monkeypatch): monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") trace_info = TraceInfo( trace_id="123", trace_location=EntityTraceLocation.from_experiment_id("0"), request_time=0, execution_duration=1, state=TraceState.OK, tags={MLFLOW_ARTIFACT_LOCATION: location}, ) repo = _get_trace_artifact_repo(trace_info) assert isinstance(repo, expected_class) assert repo.artifact_uri == expected_uri ### Prompt Registry Tests ### def test_create_prompt_as_registered_model(mock_get_request_message, mock_model_registry_store): tags = [RegisteredModelTag(key=IS_PROMPT_TAG_KEY, value="true")] mock_get_request_message.return_value = CreateRegisteredModel( name="model_1", tags=[tag.to_proto() for tag in tags] ) rm = RegisteredModel("model_1", tags=tags) mock_model_registry_store.create_registered_model.return_value = rm resp = _create_registered_model() _, args = mock_model_registry_store.create_registered_model.call_args assert args["name"] == "model_1" assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm)} def test_create_prompt_as_model_version(mock_get_request_message, mock_model_registry_store): tags = [ ModelVersionTag(key=IS_PROMPT_TAG_KEY, value="true"), ModelVersionTag(key=PROMPT_TEXT_TAG_KEY, value="some prompt text"), ] mock_get_request_message.return_value = CreateModelVersion( name="model_1", tags=[tag.to_proto() for tag in tags], source=None, run_id=None, run_link=None, ) mv = ModelVersion( name="prompt_1", version="12", creation_timestamp=123, tags=tags, run_link=None ) mock_model_registry_store.create_model_version.return_value = mv resp = _create_model_version() _, args = mock_model_registry_store.create_model_version.call_args assert args["name"] == "model_1" assert args["source"] == "" assert args["run_id"] == "" assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert args["run_link"] == "" assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)} def test_create_evaluation_dataset(mock_tracking_store, mock_evaluation_dataset): mock_tracking_store.create_dataset.return_value = mock_evaluation_dataset with app.test_request_context( method="POST", json={ "name": "test_dataset", "experiment_ids": ["0", "1"], "tags": json.dumps({"env": "test"}), }, ): _create_dataset_handler() mock_tracking_store.create_dataset.assert_called_once_with( name="test_dataset", experiment_ids=["0", "1"], tags={"env": "test"}, ) def test_get_evaluation_dataset(mock_tracking_store, mock_evaluation_dataset): mock_tracking_store.get_dataset.return_value = mock_evaluation_dataset dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context(method="GET"): _get_dataset_handler(dataset_id) mock_tracking_store.get_dataset.assert_called_once_with(dataset_id) def test_delete_evaluation_dataset(mock_tracking_store): dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context(method="DELETE"): _delete_dataset_handler(dataset_id) mock_tracking_store.delete_dataset.assert_called_once_with(dataset_id) def test_search_datasets(mock_tracking_store): from mlflow.protos.datasets_pb2 import Dataset as ProtoDataset datasets = [] for i in range(2): ds = mock.MagicMock() ds.name = f"dataset_{i}" proto = ProtoDataset() proto.dataset_id = f"d-{i:032d}" proto.name = ds.name ds.to_proto.return_value = proto datasets.append(ds) paged_list = PagedList(datasets, "next_token") mock_tracking_store.search_datasets.return_value = paged_list with app.test_request_context( method="POST", json={ "experiment_ids": ["0", "1"], "filter_string": "name = 'dataset_1'", "max_results": 10, "order_by": ["name DESC"], "page_token": "token123", }, ): _search_evaluation_datasets_handler() mock_tracking_store.search_datasets.assert_called_once_with( experiment_ids=["0", "1"], filter_string="name = 'dataset_1'", max_results=10, order_by=["name DESC"], page_token="token123", ) def test_set_dataset_tags(mock_tracking_store): dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context( method="POST", json={ "tags": json.dumps({"env": "production", "version": "2.0"}), }, ): _set_dataset_tags_handler(dataset_id) mock_tracking_store.set_dataset_tags.assert_called_once_with( dataset_id=dataset_id, tags={"env": "production", "version": "2.0"}, ) def test_delete_dataset_tag(mock_tracking_store): dataset_id = "d-1234567890abcdef1234567890abcdef" key = "deprecated_tag" with app.test_request_context(method="DELETE"): _delete_dataset_tag_handler(dataset_id, key) mock_tracking_store.delete_dataset_tag.assert_called_once_with( dataset_id=dataset_id, key=key, ) def test_upsert_dataset_records(mock_tracking_store): mock_tracking_store.upsert_dataset_records.return_value = { "inserted": 2, "updated": 0, } dataset_id = "d-1234567890abcdef1234567890abcdef" records = [ {"inputs": {"q": "test1"}, "expectations": {"score": 0.9}}, {"inputs": {"q": "test2"}, "expectations": {"score": 0.8}}, ] with app.test_request_context( method="POST", json={ "records": json.dumps(records), }, ): resp = _upsert_dataset_records_handler(dataset_id) mock_tracking_store.upsert_dataset_records.assert_called_once_with( dataset_id=dataset_id, records=records, ) response_data = json.loads(resp.get_data()) assert response_data["inserted_count"] == 2 assert response_data["updated_count"] == 0 def test_get_dataset_experiment_ids(mock_tracking_store): mock_tracking_store.get_dataset_experiment_ids.return_value = [ "exp1", "exp2", "exp3", ] dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context(method="GET"): resp = _get_dataset_experiment_ids_handler(dataset_id) mock_tracking_store.get_dataset_experiment_ids.assert_called_once_with(dataset_id=dataset_id) response_data = json.loads(resp.get_data()) assert response_data["experiment_ids"] == ["exp1", "exp2", "exp3"] def test_get_dataset_records(mock_tracking_store): records = [] for i in range(3): record = mock.MagicMock() record.dataset_id = "d-1234567890abcdef1234567890abcdef" record.dataset_record_id = f"r-00{i}" record.inputs = {"question": f"test{i}"} record.expectations = {"score": 0.9 - i * 0.1} record.tags = {} record.created_time = 1234567890 + i record.last_update_time = 1234567890 + i record.to_dict.return_value = { "dataset_id": record.dataset_id, "dataset_record_id": record.dataset_record_id, "inputs": record.inputs, "expectations": record.expectations, "tags": record.tags, "created_time": record.created_time, "last_update_time": record.last_update_time, } records.append(record) mock_tracking_store._load_dataset_records.return_value = (records, None) dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context(method="GET"): resp = _get_dataset_records_handler(dataset_id) mock_tracking_store._load_dataset_records.assert_called_with( dataset_id, max_results=1000, page_token=None ) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 3 assert records_data[0]["dataset_record_id"] == "r-000" mock_tracking_store._load_dataset_records.return_value = (records[:2], "token_page2") with app.test_request_context( method="GET", json={ "max_results": 2, "page_token": None, }, ): resp = _get_dataset_records_handler(dataset_id) mock_tracking_store._load_dataset_records.assert_called_with( dataset_id, max_results=2, page_token=None ) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 2 assert response_data["next_page_token"] == "token_page2" mock_tracking_store._load_dataset_records.return_value = (records[2:], None) with app.test_request_context( method="GET", json={ "max_results": 2, "page_token": "token_page2", }, ): resp = _get_dataset_records_handler(dataset_id) mock_tracking_store._load_dataset_records.assert_called_with( dataset_id, max_results=2, page_token="token_page2" ) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 1 assert "next_page_token" not in response_data or response_data["next_page_token"] == "" def test_get_dataset_records_empty(mock_tracking_store): mock_tracking_store._load_dataset_records.return_value = ([], None) dataset_id = "d-1234567890abcdef1234567890abcdef" with app.test_request_context(method="GET"): resp = _get_dataset_records_handler(dataset_id) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 0 assert "next_page_token" not in response_data or response_data["next_page_token"] == "" def test_get_dataset_records_pagination(mock_tracking_store): dataset_id = "d-1234567890abcdef1234567890abcdef" all_records = [] for i in range(50): record = mock.Mock() record.dataset_record_id = f"r-{i:03d}" record.inputs = {"q": f"Question {i}"} record.expectations = {"a": f"Answer {i}"} record.tags = {} record.source_type = "TRACE" record.source_id = f"trace-{i}" record.created_time = 1609459200 + i record.to_dict.return_value = { "dataset_record_id": f"r-{i:03d}", "inputs": {"q": f"Question {i}"}, "expectations": {"a": f"Answer {i}"}, "tags": {}, "source_type": "TRACE", "source_id": f"trace-{i}", "created_time": 1609459200 + i, } all_records.append(record) mock_tracking_store._load_dataset_records.return_value = (all_records[:20], "token_20") with app.test_request_context( method="GET", json={"max_results": 20}, ): resp = _get_dataset_records_handler(dataset_id) mock_tracking_store._load_dataset_records.assert_called_with( dataset_id, max_results=20, page_token=None ) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 20 assert response_data["next_page_token"] == "token_20" assert records_data[0]["dataset_record_id"] == "r-000" assert records_data[19]["dataset_record_id"] == "r-019" mock_tracking_store._load_dataset_records.return_value = (all_records[20:40], "token_40") with app.test_request_context( method="GET", json={"max_results": 20, "page_token": "token_20"}, ): resp = _get_dataset_records_handler(dataset_id) mock_tracking_store._load_dataset_records.assert_called_with( dataset_id, max_results=20, page_token="token_20" ) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 20 assert response_data["next_page_token"] == "token_40" assert records_data[0]["dataset_record_id"] == "r-020" mock_tracking_store._load_dataset_records.return_value = (all_records[40:], None) with app.test_request_context( method="GET", json={"max_results": 20, "page_token": "token_40"}, ): resp = _get_dataset_records_handler(dataset_id) response_data = json.loads(resp.get_data()) records_data = json.loads(response_data["records"]) assert len(records_data) == 10 assert "next_page_token" not in response_data or response_data["next_page_token"] == "" assert records_data[0]["dataset_record_id"] == "r-040" assert records_data[9]["dataset_record_id"] == "r-049" def test_register_scorer(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" serialized_scorer = '{"name": "accuracy_scorer"}' mock_get_request_message.return_value = RegisterScorer( experiment_id=experiment_id, name=name, serialized_scorer=serialized_scorer ) mock_scorer_version = ScorerVersion( experiment_id=experiment_id, scorer_name=name, scorer_version=1, serialized_scorer=serialized_scorer, creation_time=1234567890, scorer_id="test-scorer-id", ) mock_tracking_store.register_scorer.return_value = mock_scorer_version resp = _register_scorer() mock_tracking_store.register_scorer.assert_called_once_with( experiment_id, name, serialized_scorer ) response_data = json.loads(resp.get_data()) assert response_data == { "version": 1, "scorer_id": "test-scorer-id", "experiment_id": experiment_id, "name": name, "serialized_scorer": serialized_scorer, "creation_time": 1234567890, } def test_register_scorer_rejects_decorator_scorer(mock_get_request_message, mock_tracking_store): from mlflow.genai.scorers.scorer_utils import DECORATOR_SCORER_REGISTRATION_NOT_SUPPORTED_ERROR serialized_scorer = json.dumps({"name": "my_scorer", "call_source": " return 1.0\n"}) mock_get_request_message.return_value = RegisterScorer( experiment_id="123", name="my_scorer", serialized_scorer=serialized_scorer ) resp = _register_scorer() assert resp.status_code == 400 assert DECORATOR_SCORER_REGISTRATION_NOT_SUPPORTED_ERROR in resp.get_json()["message"] mock_tracking_store.register_scorer.assert_not_called() def test_list_scorers(mock_get_request_message, mock_tracking_store): experiment_id = "123" mock_get_request_message.return_value = ListScorers(experiment_id=experiment_id) # Create mock scorers scorers = [ ScorerVersion( experiment_id=123, scorer_name="accuracy_scorer", scorer_version=1, serialized_scorer="serialized_accuracy_scorer", creation_time=12345, ), ScorerVersion( experiment_id=123, scorer_name="safety_scorer", scorer_version=2, serialized_scorer="serialized_safety_scorer", creation_time=12345, ), ] mock_tracking_store.list_scorers.return_value = scorers resp = _list_scorers() # Verify the tracking store was called with correct arguments mock_tracking_store.list_scorers.assert_called_once_with(experiment_id) # Verify the response response_data = json.loads(resp.get_data()) assert len(response_data["scorers"]) == 2 assert response_data["scorers"][0]["scorer_name"] == "accuracy_scorer" assert response_data["scorers"][0]["scorer_version"] == 1 assert response_data["scorers"][0]["serialized_scorer"] == "serialized_accuracy_scorer" assert response_data["scorers"][1]["scorer_name"] == "safety_scorer" assert response_data["scorers"][1]["scorer_version"] == 2 assert response_data["scorers"][1]["serialized_scorer"] == "serialized_safety_scorer" def test_list_scorers_cross_experiment(mock_get_request_message, mock_tracking_store): # Empty ``experiment_id`` switches to the cross-experiment branch: the # handler walks ``search_experiments`` via ``page_token`` until exhausted, # then makes a single ``list_scorers_across_experiments`` call with the # collected ids. Mocks below force a 3-page walk so the loop is exercised # end-to-end (not just the first page). mock_get_request_message.return_value = ListScorers() def _make_page(items, token): return PagedList( [ Experiment( experiment_id=str(i), name=f"e-{i}", artifact_location="", lifecycle_stage="active", ) for i in items ], token, ) mock_tracking_store.search_experiments.side_effect = [ _make_page([1, 2], "tok-1"), _make_page([3], "tok-2"), _make_page([], None), # terminal page with no token ] mock_tracking_store.list_scorers_across_experiments.return_value = [ ScorerVersion( experiment_id=1, scorer_name="alpha", scorer_version=1, serialized_scorer="s", creation_time=0, ), ] _list_scorers() # ``search_experiments`` is called 3 times: initial + two follow-ups using # the prior page's token. ``list_scorers`` (single-experiment) must NOT be # called on this branch. assert mock_tracking_store.search_experiments.call_count == 3 page_tokens = [ c.kwargs.get("page_token") for c in mock_tracking_store.search_experiments.call_args_list ] assert page_tokens == [None, "tok-1", "tok-2"] mock_tracking_store.list_scorers.assert_not_called() mock_tracking_store.list_scorers_across_experiments.assert_called_once() # Collected experiment ids: union of every page's items, in order. call_args = mock_tracking_store.list_scorers_across_experiments.call_args assert call_args.args[0] == ["1", "2", "3"] def test_list_scorer_versions(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" mock_get_request_message.return_value = ListScorerVersions( experiment_id=experiment_id, name=name ) # Create mock scorers with multiple versions scorers = [ ScorerVersion( experiment_id=123, scorer_name="accuracy_scorer", scorer_version=1, serialized_scorer="serialized_accuracy_scorer_v1", creation_time=12345, ), ScorerVersion( experiment_id=123, scorer_name="accuracy_scorer", scorer_version=2, serialized_scorer="serialized_accuracy_scorer_v2", creation_time=12345, ), ] mock_tracking_store.list_scorer_versions.return_value = scorers resp = _list_scorer_versions() # Verify the tracking store was called with correct arguments mock_tracking_store.list_scorer_versions.assert_called_once_with(experiment_id, name) # Verify the response response_data = json.loads(resp.get_data()) assert len(response_data["scorers"]) == 2 assert response_data["scorers"][0]["scorer_version"] == 1 assert response_data["scorers"][0]["serialized_scorer"] == "serialized_accuracy_scorer_v1" assert response_data["scorers"][1]["scorer_version"] == 2 assert response_data["scorers"][1]["serialized_scorer"] == "serialized_accuracy_scorer_v2" def test_get_scorer_with_version(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" version = 2 mock_get_request_message.return_value = GetScorer( experiment_id=experiment_id, name=name, version=version ) # Mock the return value as a ScorerVersion entity mock_scorer_version = ScorerVersion( experiment_id=123, scorer_name="accuracy_scorer", scorer_version=2, serialized_scorer="serialized_accuracy_scorer_v2", creation_time=1640995200000, ) mock_tracking_store.get_scorer.return_value = mock_scorer_version resp = _get_scorer() # Verify the tracking store was called with correct arguments (positional) mock_tracking_store.get_scorer.assert_called_once_with(experiment_id, name, version) # Verify the response response_data = json.loads(resp.get_data()) assert response_data["scorer"]["experiment_id"] == 123 assert response_data["scorer"]["scorer_name"] == "accuracy_scorer" assert response_data["scorer"]["scorer_version"] == 2 assert response_data["scorer"]["serialized_scorer"] == "serialized_accuracy_scorer_v2" assert response_data["scorer"]["creation_time"] == 1640995200000 def test_get_scorer_without_version(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" mock_get_request_message.return_value = GetScorer(experiment_id=experiment_id, name=name) # Mock the return value as a ScorerVersion entity mock_scorer_version = ScorerVersion( experiment_id=123, scorer_name="accuracy_scorer", scorer_version=3, serialized_scorer="serialized_accuracy_scorer_latest", creation_time=1640995200000, ) mock_tracking_store.get_scorer.return_value = mock_scorer_version resp = _get_scorer() # Verify the tracking store was called with correct arguments (positional, version=None) mock_tracking_store.get_scorer.assert_called_once_with(experiment_id, name, None) # Verify the response response_data = json.loads(resp.get_data()) assert response_data["scorer"]["experiment_id"] == 123 assert response_data["scorer"]["scorer_name"] == "accuracy_scorer" assert response_data["scorer"]["scorer_version"] == 3 assert response_data["scorer"]["serialized_scorer"] == "serialized_accuracy_scorer_latest" assert response_data["scorer"]["creation_time"] == 1640995200000 def test_delete_scorer_with_version(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" version = 2 mock_get_request_message.return_value = DeleteScorer( experiment_id=experiment_id, name=name, version=version ) resp = _delete_scorer() # Verify the tracking store was called with correct arguments (positional) mock_tracking_store.delete_scorer.assert_called_once_with(experiment_id, name, version) # Verify the response (should be empty for delete operations) response_data = json.loads(resp.get_data()) assert response_data == {} def test_delete_scorer_without_version(mock_get_request_message, mock_tracking_store): experiment_id = "123" name = "accuracy_scorer" mock_get_request_message.return_value = DeleteScorer(experiment_id=experiment_id, name=name) resp = _delete_scorer() # Verify the tracking store was called with correct arguments (positional, version=None) mock_tracking_store.delete_scorer.assert_called_once_with(experiment_id, name, None) # Verify the response (should be empty for delete operations) response_data = json.loads(resp.get_data()) assert response_data == {} def test_get_online_scoring_configs_batch(mock_tracking_store): mock_configs = [ OnlineScoringConfig( online_scoring_config_id="cfg-1", scorer_id="scorer-1", sample_rate=0.5, filter_string="status = 'OK'", experiment_id="exp1", ), OnlineScoringConfig( online_scoring_config_id="cfg-2", scorer_id="scorer-2", sample_rate=0.8, experiment_id="exp1", ), ] mock_tracking_store.get_online_scoring_configs.return_value = mock_configs with app.test_client() as c: resp = c.get( "/ajax-api/3.0/mlflow/scorers/online-configs", query_string=[("scorer_ids", "scorer-1"), ("scorer_ids", "scorer-2")], ) assert resp.status_code == 200 data = resp.get_json() assert "configs" in data assert isinstance(data["configs"], list) assert len(data["configs"]) == 2 configs_by_id = {c["scorer_id"]: c for c in data["configs"]} assert configs_by_id["scorer-1"]["sample_rate"] == 0.5 assert configs_by_id["scorer-1"]["filter_string"] == "status = 'OK'" assert configs_by_id["scorer-2"]["sample_rate"] == 0.8 assert configs_by_id["scorer-2"].get("filter_string") is None mock_tracking_store.get_online_scoring_configs.assert_called_once_with(["scorer-1", "scorer-2"]) def test_get_online_scoring_configs_missing_param(mock_tracking_store): with app.test_client() as c: resp = c.get( "/ajax-api/3.0/mlflow/scorers/online-configs", ) assert resp.status_code == 400 data = resp.get_json() assert "scorer_ids" in data["message"] def test_calculate_trace_filter_correlation(mock_get_request_message, mock_tracking_store): experiment_ids = ["123", "456"] filter_string1 = "span.type = 'LLM'" filter_string2 = "feedback.quality > 0.8" base_filter = "request_time > 1000" mock_request = CalculateTraceFilterCorrelation( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, base_filter=base_filter, ) mock_get_request_message.return_value = mock_request mock_result = TraceFilterCorrelationResult( npmi=0.456, npmi_smoothed=0.445, filter1_count=100, filter2_count=80, joint_count=50, total_count=200, ) mock_tracking_store.calculate_trace_filter_correlation.return_value = mock_result resp = _calculate_trace_filter_correlation() mock_tracking_store.calculate_trace_filter_correlation.assert_called_once_with( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, base_filter=base_filter, ) response_data = json.loads(resp.get_data()) assert response_data["npmi"] == 0.456 assert response_data["npmi_smoothed"] == 0.445 assert response_data["filter1_count"] == 100 assert response_data["filter2_count"] == 80 assert response_data["joint_count"] == 50 assert response_data["total_count"] == 200 def test_calculate_trace_filter_correlation_without_base_filter( mock_get_request_message, mock_tracking_store ): experiment_ids = ["123"] filter_string1 = "span.type = 'LLM'" filter_string2 = "feedback.quality > 0.8" mock_request = CalculateTraceFilterCorrelation( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, ) mock_get_request_message.return_value = mock_request mock_result = TraceFilterCorrelationResult( npmi=0.789, npmi_smoothed=0.775, filter1_count=50, filter2_count=40, joint_count=30, total_count=100, ) mock_tracking_store.calculate_trace_filter_correlation.return_value = mock_result resp = _calculate_trace_filter_correlation() mock_tracking_store.calculate_trace_filter_correlation.assert_called_once_with( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, base_filter=None, ) response_data = json.loads(resp.get_data()) assert response_data["npmi"] == 0.789 assert response_data["npmi_smoothed"] == 0.775 assert response_data["filter1_count"] == 50 assert response_data["filter2_count"] == 40 assert response_data["joint_count"] == 30 assert response_data["total_count"] == 100 def test_calculate_trace_filter_correlation_with_nan_npmi( mock_get_request_message, mock_tracking_store ): experiment_ids = ["123"] filter_string1 = "span.type = 'LLM'" filter_string2 = "feedback.quality > 0.8" mock_request = CalculateTraceFilterCorrelation( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, ) mock_get_request_message.return_value = mock_request mock_result = TraceFilterCorrelationResult( npmi=float("nan"), npmi_smoothed=None, filter1_count=0, filter2_count=0, joint_count=0, total_count=100, ) mock_tracking_store.calculate_trace_filter_correlation.return_value = mock_result resp = _calculate_trace_filter_correlation() mock_tracking_store.calculate_trace_filter_correlation.assert_called_once_with( experiment_ids=experiment_ids, filter_string1=filter_string1, filter_string2=filter_string2, base_filter=None, ) response_data = json.loads(resp.get_data()) assert "npmi" not in response_data assert "npmi_smoothed" not in response_data assert response_data["filter1_count"] == 0 assert response_data["filter2_count"] == 0 assert response_data["joint_count"] == 0 assert response_data["total_count"] == 100 def test_databricks_tracking_store_registration(): registry = TrackingStoreRegistryWrapper() # Test that the correct store type is returned for databricks scheme store = registry.get_store("databricks", artifact_uri=None) assert isinstance(store, DatabricksTracingRestStore) # Verify that the store was created with the right get_host_creds function # The RestStore should have a get_host_creds attribute that is a partial function assert hasattr(store, "get_host_creds") assert store.get_host_creds.func.__name__ == "get_databricks_host_creds" assert store.get_host_creds.args == ("databricks",) def test_databricks_model_registry_store_registration(): registry = ModelRegistryStoreRegistryWrapper() # Test that the correct store type is returned for databricks store = registry.get_store("databricks") assert isinstance(store, ModelRegistryRestStore) # Verify that the store was created with the right get_host_creds function assert hasattr(store, "get_host_creds") assert store.get_host_creds.func.__name__ == "get_databricks_host_creds" assert store.get_host_creds.args == ("databricks",) # Test that the correct store type is returned for databricks-uc uc_store = registry.get_store("databricks-uc") assert isinstance(uc_store, UcModelRegistryStore) # Verify that the UC store was created with the right get_host_creds function # Note: UcModelRegistryStore uses get_databricks_host_creds internally, # not get_databricks_uc_host_creds assert hasattr(uc_store, "get_host_creds") assert uc_store.get_host_creds.func.__name__ == "get_databricks_host_creds" assert uc_store.get_host_creds.args == ("databricks-uc",) # Also verify it has tracking_uri set assert hasattr(uc_store, "tracking_uri") # The tracking_uri will be set based on environment/test config # In test environment, it may be set to a test sqlite database assert uc_store.tracking_uri is not None def test_search_experiments_empty_page_token(mock_get_request_message, mock_tracking_store): # Create proto without setting page_token - it defaults to empty string search_experiments_proto = SearchExperiments() search_experiments_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_experiments_proto.page_token == "" mock_get_request_message.return_value = search_experiments_proto mock_tracking_store.search_experiments.return_value = PagedList([], None) _search_experiments() # Verify that search_experiments was called with page_token=None (not empty string) mock_tracking_store.search_experiments.assert_called_once() call_kwargs = mock_tracking_store.search_experiments.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_search_registered_models_empty_page_token( mock_get_request_message, mock_model_registry_store ): # Create proto without setting page_token - it defaults to empty string search_registered_models_proto = SearchRegisteredModels() search_registered_models_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_registered_models_proto.page_token == "" mock_get_request_message.return_value = search_registered_models_proto mock_model_registry_store.search_registered_models.return_value = PagedList([], None) _search_registered_models() # Verify that search_registered_models was called with page_token=None (not empty string) mock_model_registry_store.search_registered_models.assert_called_once() call_kwargs = mock_model_registry_store.search_registered_models.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_search_model_versions_empty_page_token( mock_get_request_message, mock_model_registry_store ): # Create proto without setting page_token - it defaults to empty string search_model_versions_proto = SearchModelVersions() search_model_versions_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_model_versions_proto.page_token == "" mock_get_request_message.return_value = search_model_versions_proto mock_model_registry_store.search_model_versions.return_value = PagedList([], None) _search_model_versions() # Verify that search_model_versions was called with page_token=None (not empty string) mock_model_registry_store.search_model_versions.assert_called_once() call_kwargs = mock_model_registry_store.search_model_versions.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_search_traces_v3_empty_page_token(mock_get_request_message, mock_tracking_store): # Create proto without setting page_token - it defaults to empty string # SearchTracesV3 requires locations field search_traces_proto = SearchTracesV3() location = TraceLocation() location.mlflow_experiment.experiment_id = "1" search_traces_proto.locations.append(location) search_traces_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_traces_proto.page_token == "" mock_get_request_message.return_value = search_traces_proto mock_tracking_store.search_traces.return_value = ([], None) _search_traces_v3() # Verify that search_traces was called with page_token=None (not empty string) mock_tracking_store.search_traces.assert_called_once() call_kwargs = mock_tracking_store.search_traces.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_deprecated_search_traces_v2_empty_page_token( mock_get_request_message, mock_tracking_store ): # Create proto without setting page_token - it defaults to empty string search_traces_proto = SearchTraces() search_traces_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_traces_proto.page_token == "" mock_get_request_message.return_value = search_traces_proto mock_tracking_store.search_traces.return_value = ([], None) _deprecated_search_traces_v2() # Verify that search_traces was called with page_token=None (not empty string) mock_tracking_store.search_traces.assert_called_once() call_kwargs = mock_tracking_store.search_traces.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_search_logged_models_empty_page_token(mock_get_request_message, mock_tracking_store): # Create proto without setting page_token - it defaults to empty string search_logged_models_proto = SearchLoggedModels() search_logged_models_proto.max_results = 10 # Verify that proto's default page_token is empty string assert search_logged_models_proto.page_token == "" mock_get_request_message.return_value = search_logged_models_proto mock_tracking_store.search_logged_models.return_value = PagedList([], None) _search_logged_models() # Verify that search_logged_models was called with page_token=None (not empty string) mock_tracking_store.search_logged_models.assert_called_once() call_kwargs = mock_tracking_store.search_logged_models.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_list_webhooks_empty_page_token(mock_get_request_message, mock_model_registry_store): # Create proto without setting page_token - it defaults to empty string list_webhooks_proto = ListWebhooks() list_webhooks_proto.max_results = 10 # Verify that proto's default page_token is empty string assert list_webhooks_proto.page_token == "" mock_get_request_message.return_value = list_webhooks_proto mock_model_registry_store.list_webhooks.return_value = PagedList([], None) _list_webhooks() # Verify that list_webhooks was called with page_token=None (not empty string) mock_model_registry_store.list_webhooks.assert_called_once() call_kwargs = mock_model_registry_store.list_webhooks.call_args.kwargs assert call_kwargs.get("page_token") is None assert call_kwargs.get("max_results") == 10 def test_batch_get_traces_handler(mock_get_request_message, mock_tracking_store): trace_id_1 = "test-trace-123" trace_id_2 = "test-trace-456" get_traces_proto = BatchGetTraces(trace_ids=[trace_id_1, trace_id_2]) mock_get_request_message.return_value = get_traces_proto otel_span = OTelReadableSpan( name="test", context=build_otel_context(123, 234), parent=None, start_time=100, end_time=200, attributes={ "mlflow.spanInputs": json.dumps("inputs"), "mlflow.spanOutputs": json.dumps("outputs"), "mlflow.spanType": json.dumps("span_type"), }, ) mock_span = Span(otel_span) # Create mock traces to return mock_trace_1 = Trace( info=TraceInfo( trace_id=trace_id_1, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=5000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) mock_trace_2 = Trace( info=TraceInfo( trace_id=trace_id_2, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=3000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) mock_tracking_store.batch_get_traces.return_value = [mock_trace_1, mock_trace_2] # Call the handler response = _batch_get_traces() # Verify the store was called with the correct trace IDs mock_tracking_store.batch_get_traces.assert_called_once_with([trace_id_1, trace_id_2], None) # Verify response was created assert response is not None assert response.status_code == 200 traces = json.loads(response.get_data())["traces"] assert len(traces) == 2 assert len(traces[0]["spans"]) == 1 assert len(traces[1]["spans"]) == 1 def test_batch_get_traces_handler_empty_list(mock_get_request_message, mock_tracking_store): get_traces_proto = BatchGetTraces() mock_get_request_message.return_value = get_traces_proto mock_tracking_store.batch_get_traces.return_value = [] response = _batch_get_traces() mock_tracking_store.batch_get_traces.assert_called_once_with([], None) # Verify response was created assert response is not None assert response.status_code == 200 def test_batch_get_trace_infos_handler(mock_get_request_message, mock_tracking_store): trace_id_1 = "test-trace-123" trace_id_2 = "test-trace-456" mock_get_request_message.return_value = BatchGetTraceInfos(trace_ids=[trace_id_1, trace_id_2]) mock_trace_info_1 = TraceInfo( trace_id=trace_id_1, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=5000, state=TraceState.OK, ) mock_trace_info_2 = TraceInfo( trace_id=trace_id_2, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=3000, state=TraceState.OK, ) mock_tracking_store.batch_get_trace_infos.return_value = [ mock_trace_info_1, mock_trace_info_2, ] response = _batch_get_trace_infos() mock_tracking_store.batch_get_trace_infos.assert_called_once_with([trace_id_1, trace_id_2]) assert response is not None assert response.status_code == 200 trace_infos = json.loads(response.get_data())["trace_infos"] assert len(trace_infos) == 2 assert trace_infos[0]["trace_id"] == trace_id_1 assert trace_infos[1]["trace_id"] == trace_id_2 def test_get_trace_handler(mock_get_request_message, mock_tracking_store): trace_id = "test-trace-123" get_trace_proto = GetTrace(trace_id=trace_id, allow_partial=True) mock_get_request_message.return_value = get_trace_proto otel_span = OTelReadableSpan( name="test", context=build_otel_context(123, 234), parent=None, start_time=100, end_time=200, attributes={ "mlflow.spanInputs": json.dumps("inputs"), "mlflow.spanOutputs": json.dumps("outputs"), "mlflow.spanType": json.dumps("span_type"), }, ) mock_span = Span(otel_span) mock_trace = Trace( info=TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=5000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) mock_tracking_store.get_trace.return_value = mock_trace response = _get_trace() mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert "trace" in response_data trace = response_data["trace"] assert trace["trace_info"]["trace_id"] == trace_id assert len(trace["spans"]) == 1 def test_get_trace_handler_with_allow_partial_false(mock_get_request_message, mock_tracking_store): trace_id = "test-trace-456" get_trace_proto = GetTrace(trace_id=trace_id, allow_partial=False) mock_get_request_message.return_value = get_trace_proto otel_span = OTelReadableSpan( name="test", context=build_otel_context(123, 234), parent=None, start_time=100, end_time=200, attributes={}, ) mock_span = Span(otel_span) mock_trace = Trace( info=TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=5000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) mock_tracking_store.get_trace.return_value = mock_trace response = _get_trace() mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=False) assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert "trace" in response_data def test_get_trace_handler_not_found(mock_get_request_message, mock_tracking_store): trace_id = "non-existent-trace" get_trace_proto = GetTrace(trace_id=trace_id) mock_get_request_message.return_value = get_trace_proto mock_tracking_store.get_trace.side_effect = MlflowException( f"Trace with ID {trace_id} is not found.", error_code=RESOURCE_DOES_NOT_EXIST, ) response = _get_trace() mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=False) assert response is not None assert response.status_code == 404 response_data = json.loads(response.get_data()) assert "error_code" in response_data assert response_data["error_code"] == "RESOURCE_DOES_NOT_EXIST" def test_get_trace_artifact_handler(mock_tracking_store): trace_id = "test-trace-artifact-123" otel_span = OTelReadableSpan( name="test_span", context=build_otel_context(123, 234), parent=None, start_time=100, end_time=200, attributes={ "mlflow.spanInputs": json.dumps({"input": "test_input"}), "mlflow.spanOutputs": json.dumps({"output": "test_output"}), }, ) mock_span = Span(otel_span) mock_trace = Trace( info=TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("1"), request_time=1234567890, execution_duration=5000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) mock_tracking_store.get_trace.return_value = mock_trace mock_tracking_store.batch_get_traces.return_value = [mock_trace] with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() # Verify the store was called correctly mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) # Verify response headers and status assert response is not None assert response.status_code == 200 assert response.headers["Content-Disposition"] == "attachment; filename=traces.json" def test_get_trace_artifact_handler_missing_request_id(mock_tracking_store): with app.test_request_context(method="GET"): response = get_trace_artifact_handler() assert response.status_code == 400 response_data = json.loads(response.get_data()) assert "error_code" in response_data assert response_data["error_code"] == "BAD_REQUEST" assert 'must include the "request_id" query parameter' in response_data["message"] def test_get_trace_artifact_handler_trace_not_found(mock_tracking_store): trace_id = "non-existent-trace" mock_tracking_store.get_trace.side_effect = MlflowException( f"Trace with ID {trace_id} is not found.", error_code=RESOURCE_DOES_NOT_EXIST, ) with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) assert response.status_code == 404 response_data = json.loads(response.get_data()) assert "error_code" in response_data assert response_data["error_code"] == "RESOURCE_DOES_NOT_EXIST" assert f"Trace with ID {trace_id} is not found" in response_data["message"] def test_get_trace_artifact_handler_fallback_to_batch_get_traces(mock_tracking_store): trace_id = "test-trace-batch-123" otel_span = OTelReadableSpan( name="test_span_batch", context=build_otel_context(456, 789), parent=None, start_time=100, end_time=200, attributes={ "mlflow.spanInputs": json.dumps({"input": "batch_input"}), "mlflow.spanOutputs": json.dumps({"output": "batch_output"}), }, ) mock_span = Span(otel_span) mock_trace = Trace( info=TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("2"), request_time=1234567890, execution_duration=3000, state=TraceState.OK, ), data=TraceData(spans=[mock_span]), ) # Simulate get_trace not being implemented mock_tracking_store.get_trace.side_effect = MlflowNotImplementedException( "get_trace is not implemented" ) mock_tracking_store.batch_get_traces.return_value = [mock_trace] with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() # Verify both methods were called mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) mock_tracking_store.batch_get_traces.assert_called_once_with([trace_id], None) # Verify successful response assert response is not None assert response.status_code == 200 assert response.headers["Content-Disposition"] == "attachment; filename=traces.json" def test_get_trace_artifact_handler_batch_get_traces_not_found(mock_tracking_store): trace_id = "non-existent-batch-trace" # Simulate get_trace not being implemented mock_tracking_store.get_trace.side_effect = MlflowNotImplementedException( "get_trace is not implemented" ) # batch_get_traces returns empty list (trace not found) mock_tracking_store.batch_get_traces.return_value = [] with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() # Verify both methods were called mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) mock_tracking_store.batch_get_traces.assert_called_once_with([trace_id], None) # Verify 404 response assert response.status_code == 404 response_data = json.loads(response.get_data()) assert "error_code" in response_data assert response_data["error_code"] == "RESOURCE_DOES_NOT_EXIST" assert f"Trace with id={trace_id} not found" in response_data["message"] def test_get_trace_artifact_handler_fallback_to_artifact_repo(mock_tracking_store): trace_id = "test-trace-artifact-repo-123" trace_info = TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("3"), request_time=1234567890, execution_duration=4000, state=TraceState.OK, ) trace_data = { "spans": [ { "name": "artifact_span", "context": {"trace_id": trace_id, "span_id": "123"}, "parent_id": None, "start_time": 100, "end_time": 200, "status_code": "OK", "status_message": "", "attributes": {}, "events": [], } ] } # Simulate batch_get_traces not being implemented mock_tracking_store.get_trace.side_effect = MlflowNotImplementedException( "get_trace is not implemented" ) mock_tracking_store.batch_get_traces.side_effect = MlflowNotImplementedException( "batch_get_traces is not implemented" ) mock_tracking_store.get_trace_info.return_value = trace_info # Mock the artifact repo mock_artifact_repo = mock.MagicMock() mock_artifact_repo.download_trace_data.return_value = trace_data with mock.patch( "mlflow.server.handlers._get_trace_artifact_repo", return_value=mock_artifact_repo ): with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() # Verify the fallback path was taken mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) mock_tracking_store.batch_get_traces.assert_called_once_with([trace_id], None) mock_tracking_store.get_trace_info.assert_called_once_with(trace_id) mock_artifact_repo.download_trace_data.assert_called_once() # Verify successful response assert response is not None assert response.status_code == 200 assert response.headers["Content-Disposition"] == "attachment; filename=traces.json" def test_get_trace_artifact_handler_with_attachment_path(mock_tracking_store): trace_id = "tr-test-attachment-123" attachment_id = "a1b2c3d4-e5f6-4890-abcd-ef1234567890" trace_info = TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("3"), request_time=1234567890, execution_duration=4000, state=TraceState.OK, ) mock_tracking_store.get_trace_info.return_value = trace_info mock_artifact_repo = mock.MagicMock() mock_artifact_repo.download_trace_attachment.return_value = b"\x89PNG fake image" with mock.patch( "mlflow.server.handlers._get_trace_artifact_repo", return_value=mock_artifact_repo ): query = {"request_id": trace_id, "path": attachment_id} with app.test_request_context(method="GET", query_string=query): response = get_trace_artifact_handler() mock_tracking_store.get_trace_info.assert_called_once_with(trace_id) mock_artifact_repo.download_trace_attachment.assert_called_once_with(attachment_id) assert response.status_code == 200 assert response.headers["Content-Type"] == "application/octet-stream" assert response.headers["Content-Disposition"] == f"attachment; filename={attachment_id}" assert response.headers["X-Content-Type-Options"] == "nosniff" def test_get_trace_artifact_handler_falls_back_to_archive_repo(mock_tracking_store): trace_id = "tr-test-archive-fallback" trace_info = TraceInfo( trace_id=trace_id, trace_location=EntityTraceLocation.from_experiment_id("3"), request_time=1234567890, execution_duration=4000, state=TraceState.OK, tags={ MLFLOW_ARTIFACT_LOCATION: "dbfs:/trace-artifacts", TraceTagKey.SPANS_LOCATION: SpansLocation.ARCHIVE_REPO.value, TraceTagKey.ARCHIVE_LOCATION: "dbfs:/trace-archive", }, ) mock_tracking_store.get_trace.side_effect = MlflowTracingException("archive-backed trace") mock_tracking_store.get_trace_info.return_value = trace_info mock_archive_repo = mock.MagicMock() mock_archive_repo.download_archived_trace_data.return_value = TraceData(spans=[]) with mock.patch( "mlflow.server.handlers._get_trace_archive_repo", return_value=mock_archive_repo ): with app.test_request_context(method="GET", query_string={"request_id": trace_id}): response = get_trace_artifact_handler() mock_tracking_store.get_trace.assert_called_once_with(trace_id, allow_partial=True) mock_tracking_store.get_trace_info.assert_called_once_with(trace_id) mock_archive_repo.download_archived_trace_data.assert_called_once() assert response.status_code == 200 assert response.headers["Content-Disposition"] == "attachment; filename=traces.json" def test_get_trace_artifact_handler_attachment_missing_request_id(): query = {"path": "a1b2c3d4-e5f6-4890-abcd-ef1234567890"} with app.test_request_context(method="GET", query_string=query): response = get_trace_artifact_handler() assert response.status_code == 400 def test_get_trace_artifact_handler_attachment_trace_not_found(mock_tracking_store): mock_tracking_store.get_trace_info.return_value = None query = {"request_id": "tr-nonexistent", "path": "a1b2c3d4-e5f6-4890-abcd-ef1234567890"} with app.test_request_context(method="GET", query_string=query): response = get_trace_artifact_handler() assert response.status_code == 404 def test_delete_trace_tag_v2_handler(mock_get_request_message, mock_tracking_store): """Test v2 delete_trace_tag handler with request_id parameter. Verifies that when the Flask route uses request_id path parameter, the _delete_trace_tag handler is called and invokes store.delete_trace_tag(). """ request_id = "tr-123v2" tag_key = "tk" # Create the request message request_msg = DeleteTraceTag(key=tag_key) mock_get_request_message.return_value = request_msg # Call the v2 handler with request_id parameter response = _delete_trace_tag(request_id=request_id) # Verify the store method was called with correct parameters mock_tracking_store.delete_trace_tag.assert_called_once_with(request_id, tag_key) assert response is not None assert response.status_code == 200 def test_delete_trace_tag_v3_handler(mock_get_request_message, mock_tracking_store): """Test v3 delete_trace_tag handler with trace_id parameter. Verifies that when the Flask route uses trace_id path parameter, the _delete_trace_tag_v3 handler is called and invokes store.delete_trace_tag(). This is similar to v2 but uses the v3 proto message and route parameter naming. """ trace_id = "tr-v3-456" tag_key = "tk" # Create the request message with V3 request_msg = DeleteTraceTagV3(key=tag_key) mock_get_request_message.return_value = request_msg # Call the v3 handler with trace_id parameter response = _delete_trace_tag_v3(trace_id=trace_id) # Verify the store method was called with correct parameters # Both v2 and v3 call the same store method mock_tracking_store.delete_trace_tag.assert_called_once_with(trace_id, tag_key) assert response is not None assert response.status_code == 200 @pytest.mark.parametrize( "tag_key", [ TraceTagKey.SPANS_LOCATION, TraceTagKey.ARCHIVE_LOCATION, ], ) def test_delete_trace_tag_v2_handler_rejects_archival_immutable_tags( mock_get_request_message, mock_tracking_store, tag_key ): request_msg = DeleteTraceTag(key=tag_key) mock_get_request_message.return_value = request_msg response = _delete_trace_tag(request_id="tr-123v2") assert response.status_code == 400 response_data = json.loads(response.get_data()) assert response_data["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert response_data["message"] == ( f"Tag '{tag_key}' is immutable and cannot be deleted on a trace." ) mock_tracking_store.delete_trace_tag.assert_not_called() @pytest.mark.parametrize( "tag_key", [ TraceTagKey.SPANS_LOCATION, TraceTagKey.ARCHIVE_LOCATION, ], ) def test_delete_trace_tag_v3_handler_rejects_archival_immutable_tags( mock_get_request_message, mock_tracking_store, tag_key ): request_msg = DeleteTraceTagV3(key=tag_key) mock_get_request_message.return_value = request_msg response = _delete_trace_tag_v3(trace_id="tr-v3-456") assert response.status_code == 400 response_data = json.loads(response.get_data()) assert response_data["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert response_data["message"] == ( f"Tag '{tag_key}' is immutable and cannot be deleted on a trace." ) mock_tracking_store.delete_trace_tag.assert_not_called() def test_delete_trace_tag_v2_handler_allows_clearing_archival_failure( mock_get_request_message, mock_tracking_store ): request_msg = DeleteTraceTag(key=TraceTagKey.ARCHIVAL_FAILURE) mock_get_request_message.return_value = request_msg response = _delete_trace_tag(request_id="tr-123v2") assert response.status_code == 200 mock_tracking_store.delete_trace_tag.assert_called_once_with( "tr-123v2", TraceTagKey.ARCHIVAL_FAILURE ) def test_delete_trace_tag_v3_handler_allows_clearing_archival_failure( mock_get_request_message, mock_tracking_store ): request_msg = DeleteTraceTagV3(key=TraceTagKey.ARCHIVAL_FAILURE) mock_get_request_message.return_value = request_msg response = _delete_trace_tag_v3(trace_id="tr-v3-456") assert response.status_code == 200 mock_tracking_store.delete_trace_tag.assert_called_once_with( "tr-v3-456", TraceTagKey.ARCHIVAL_FAILURE ) def test_set_trace_tag_v2_handler(mock_get_request_message, mock_tracking_store): """Test v2 set_trace_tag handler with request_id parameter. Verifies that when the Flask route uses request_id path parameter, the _set_trace_tag handler is called and invokes store.set_trace_tag(). """ trace_id = "tr-test-v2-123" tag_key = "tk" tag_value = "tv" # Create the request message request_msg = SetTraceTag(key=tag_key, value=tag_value) mock_get_request_message.return_value = request_msg # Call the v2 handler with request_id parameter response = _set_trace_tag(request_id=trace_id) # Verify the store method was called with correct parameters mock_tracking_store.set_trace_tag.assert_called_once_with(trace_id, tag_key, tag_value) # Verify response was created (200 status) assert response is not None assert response.status_code == 200 def test_set_trace_tag_v3_handler(mock_get_request_message, mock_tracking_store): """Test v3 set_trace_tag handler with trace_id parameter. Verifies that when the Flask route uses trace_id path parameter, the _set_trace_tag_v3 handler is called and invokes store.set_trace_tag(). This is similar to v2 but uses the v3 proto message and route parameter naming. """ trace_id = "tr-test-v3-456" tag_key = "tk" tag_value = "tv" # Create the request message (v3 version) request_msg = SetTraceTagV3(key=tag_key, value=tag_value) mock_get_request_message.return_value = request_msg # Call the v3 handler with trace_id parameter response = _set_trace_tag_v3(trace_id=trace_id) # Verify the store method was called with correct parameters # Note: Both handlers call the same store method mock_tracking_store.set_trace_tag.assert_called_once_with(trace_id, tag_key, tag_value) # Verify response was created (200 status) assert response is not None assert response.status_code == 200 @pytest.mark.parametrize( "tag_key", [ TraceTagKey.SPANS_LOCATION, TraceTagKey.ARCHIVE_LOCATION, TraceTagKey.ARCHIVAL_FAILURE, ], ) def test_set_trace_tag_v2_handler_rejects_archival_immutable_tags( mock_get_request_message, mock_tracking_store, tag_key ): request_msg = SetTraceTag(key=tag_key, value="tv") mock_get_request_message.return_value = request_msg response = _set_trace_tag(request_id="tr-test-v2-123") assert response.status_code == 400 response_data = json.loads(response.get_data()) assert response_data["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert response_data["message"] == f"Tag '{tag_key}' is immutable and cannot be set on a trace." mock_tracking_store.set_trace_tag.assert_not_called() @pytest.mark.parametrize( "tag_key", [ TraceTagKey.SPANS_LOCATION, TraceTagKey.ARCHIVE_LOCATION, TraceTagKey.ARCHIVAL_FAILURE, ], ) def test_set_trace_tag_v3_handler_rejects_archival_immutable_tags( mock_get_request_message, mock_tracking_store, tag_key ): request_msg = SetTraceTagV3(key=tag_key, value="tv") mock_get_request_message.return_value = request_msg response = _set_trace_tag_v3(trace_id="tr-test-v3-456") assert response.status_code == 400 response_data = json.loads(response.get_data()) assert response_data["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert response_data["message"] == f"Tag '{tag_key}' is immutable and cannot be set on a trace." mock_tracking_store.set_trace_tag.assert_not_called() def test_link_prompts_to_trace_handler(mock_get_request_message, mock_tracking_store): """Test link_prompts_to_trace handler. Verifies that the handler correctly parses the request and calls store.link_prompts_to_trace() with the appropriate PromptVersion objects. """ trace_id = "tr-test-123" prompt_versions_refs = [ LinkPromptsToTrace.PromptVersionRef(name="prompt1", version="1"), LinkPromptsToTrace.PromptVersionRef(name="prompt2", version="2"), ] # Create the request message request_msg = LinkPromptsToTrace(trace_id=trace_id, prompt_versions=prompt_versions_refs) mock_get_request_message.return_value = request_msg # Call the handler response = _link_prompts_to_trace() # Verify the store method was called with correct parameters # The handler should convert PromptVersionRef to PromptVersion objects call_args = mock_tracking_store.link_prompts_to_trace.call_args assert call_args[1]["trace_id"] == trace_id prompt_versions = call_args[1]["prompt_versions"] assert len(prompt_versions) == 2 assert isinstance(prompt_versions[0], PromptVersion) assert prompt_versions[0].name == "prompt1" assert prompt_versions[0].version == 1 assert isinstance(prompt_versions[1], PromptVersion) assert prompt_versions[1].name == "prompt2" assert prompt_versions[1].version == 2 # Verify response was created (200 status) assert response is not None assert response.status_code == 200 def test_list_providers(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/supported-providers") assert response.status_code == 200 data = response.get_json() assert "providers" in data assert isinstance(data["providers"], list) assert len(data["providers"]) > 0 assert "openai" in data["providers"] def test_list_providers_with_allowed_filter(monkeypatch): monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "openai,anthropic") with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/supported-providers") assert response.status_code == 200 data = response.get_json() assert "openai" in data["providers"] assert "anthropic" in data["providers"] assert "gemini" not in data["providers"] assert "bedrock" not in data["providers"] def test_list_models(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/supported-models?provider=openai") assert response.status_code == 200 data = response.get_json() assert "models" in data assert isinstance(data["models"], list) assert len(data["models"]) > 0 def test_list_models_all_providers(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/supported-models") assert response.status_code == 200 data = response.get_json() assert "models" in data assert isinstance(data["models"], list) assert len(data["models"]) > 0 def test_get_provider_config(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/provider-config?provider=openai") assert response.status_code == 200 data = response.get_json() assert "auth_modes" in data assert "default_mode" in data assert data["default_mode"] == "api_key" assert len(data["auth_modes"]) >= 1 api_key_mode = data["auth_modes"][0] assert api_key_mode["mode"] == "api_key" def test_get_provider_config_with_multiple_auth_modes(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/provider-config?provider=bedrock") assert response.status_code == 200 data = response.get_json() assert "auth_modes" in data assert data["default_mode"] == "api_key" assert len(data["auth_modes"]) >= 2 access_keys_mode = next(m for m in data["auth_modes"] if m["mode"] == "access_keys") assert len(access_keys_mode["secret_fields"]) == 2 assert any(f["name"] == "aws_secret_access_key" for f in access_keys_mode["secret_fields"]) assert any(f["name"] == "aws_region_name" for f in access_keys_mode["config_fields"]) iam_role_mode = next(m for m in data["auth_modes"] if m["mode"] == "iam_role") assert any(f["name"] == "aws_role_name" for f in iam_role_mode["config_fields"]) def test_get_provider_config_missing_provider(): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/provider-config") assert response.status_code == 400 def test_get_provider_config_with_allowed_filter(monkeypatch): monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "anthropic") with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/gateway/provider-config?provider=openai") assert response.status_code == 400 data = response.get_json() assert "not allowed" in data["message"] response = c.get("/ajax-api/3.0/mlflow/gateway/provider-config?provider=anthropic") assert response.status_code == 200 @pytest.mark.parametrize( "invalid_name", [ "invalid name", # space "invalid/name", # slash "invalid?name", # question mark "invalid&name", # ampersand "invalid#name", # hash "invalid@name", # at sign "invalid:name", # colon "日本語", # unicode (Japanese) "naïve", # unicode (accented) ], ) def test_create_gateway_endpoint_rejects_invalid_name(mock_get_request_message, invalid_name): from mlflow.protos.service_pb2 import CreateGatewayEndpoint from mlflow.server.handlers import _create_gateway_endpoint request_msg = CreateGatewayEndpoint() request_msg.name = invalid_name mock_get_request_message.return_value = request_msg response = _create_gateway_endpoint() assert response.status_code == 400 response_data = json.loads(response.get_data()) assert "Invalid endpoint name" in response_data["message"] assert response_data["error_code"] == "INVALID_PARAMETER_VALUE" @pytest.mark.parametrize( "invalid_name", [ "invalid name", # space "invalid/name", # slash "invalid?name", # question mark "invalid&name", # ampersand ], ) def test_update_gateway_endpoint_rejects_invalid_name(mock_get_request_message, invalid_name): from mlflow.protos.service_pb2 import UpdateGatewayEndpoint from mlflow.server.handlers import _update_gateway_endpoint request_msg = UpdateGatewayEndpoint() request_msg.endpoint_id = "test-endpoint-id" request_msg.name = invalid_name mock_get_request_message.return_value = request_msg response = _update_gateway_endpoint() assert response.status_code == 400 response_data = json.loads(response.get_data()) assert "Invalid endpoint name" in response_data["message"] assert response_data["error_code"] == "INVALID_PARAMETER_VALUE" def test_get_gateway_endpoint_by_endpoint_id(mock_get_request_message, mock_tracking_store): request_msg = GetGatewayEndpoint() request_msg.endpoint_id = "ep-123" mock_get_request_message.return_value = request_msg mock_endpoint = mock.MagicMock() mock_endpoint.to_proto.return_value = GatewayEndpoint(endpoint_id="ep-123") mock_tracking_store.get_gateway_endpoint.return_value = mock_endpoint response = _get_gateway_endpoint() mock_tracking_store.get_gateway_endpoint.assert_called_once_with( endpoint_id="ep-123", name=None ) assert response.status_code == 200 def test_get_gateway_endpoint_by_name(mock_get_request_message, mock_tracking_store): request_msg = GetGatewayEndpoint() request_msg.name = "my-endpoint" mock_get_request_message.return_value = request_msg mock_endpoint = mock.MagicMock() mock_endpoint.to_proto.return_value = GatewayEndpoint(endpoint_id="ep-456", name="my-endpoint") mock_tracking_store.get_gateway_endpoint.return_value = mock_endpoint response = _get_gateway_endpoint() mock_tracking_store.get_gateway_endpoint.assert_called_once_with( endpoint_id=None, name="my-endpoint" ) assert response.status_code == 200 def test_query_trace_metrics_handler(mock_get_request_message, mock_tracking_store): experiment_ids = ["exp1", "exp2"] metric_name = "latency" # Create aggregation protos aggregations_proto = [ MetricAggregation(aggregation_type=AggregationType.AVG).to_proto(), MetricAggregation( aggregation_type=AggregationType.PERCENTILE, percentile_value=95.0 ).to_proto(), ] # Create the request message request_msg = QueryTraceMetrics( experiment_ids=experiment_ids, view_type=MetricViewType.TRACES.to_proto(), metric_name=metric_name, aggregations=aggregations_proto, dimensions=["status", "model"], filters=["status = 'OK'"], time_interval_seconds=3600, start_time_ms=1000000, end_time_ms=2000000, max_results=100, page_token="token123", ) mock_get_request_message.return_value = request_msg # Create mock result mock_data_points = [ MetricDataPoint( metric_name="latency", dimensions={"status": "OK", "model": "gpt-4"}, values={"AVG": 150.5, "P95.0": 200.0}, ), MetricDataPoint( metric_name="latency", dimensions={"status": "ERROR", "model": "gpt-4"}, values={"AVG": 50.0, "P95.0": 75.0}, ), ] # Create a mock result object with next_page_token attribute mock_result = mock.MagicMock() mock_result.__iter__ = mock.MagicMock(return_value=iter(mock_data_points)) mock_result.token = "next_token" mock_tracking_store.query_trace_metrics.return_value = mock_result # Call the handler response = _query_trace_metrics() mock_tracking_store.query_trace_metrics.assert_called_once_with( experiment_ids=experiment_ids, view_type=MetricViewType.TRACES, metric_name=metric_name, aggregations=[ MetricAggregation(aggregation_type=AggregationType.AVG), MetricAggregation(aggregation_type=AggregationType.PERCENTILE, percentile_value=95.0), ], dimensions=["status", "model"], filters=["status = 'OK'"], time_interval_seconds=3600, start_time_ms=1000000, end_time_ms=2000000, max_results=100, page_token="token123", ) assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert "data_points" in response_data assert len(response_data["data_points"]) == 2 assert response_data["data_points"][0] == asdict(mock_data_points[0]) assert response_data["data_points"][1] == asdict(mock_data_points[1]) assert response_data["next_page_token"] == "next_token" def test_query_trace_metrics_handler_empty_result(mock_get_request_message, mock_tracking_store): request_msg = QueryTraceMetrics( experiment_ids=["exp1"], view_type=MetricViewType.TRACES.to_proto(), metric_name="latency", aggregations=[MetricAggregation(aggregation_type=AggregationType.AVG).to_proto()], ) mock_get_request_message.return_value = request_msg mock_result = mock.MagicMock() mock_result.__iter__ = mock.MagicMock(return_value=iter([])) mock_result.token = None mock_tracking_store.query_trace_metrics.return_value = mock_result response = _query_trace_metrics() mock_tracking_store.query_trace_metrics.assert_called_once_with( experiment_ids=["exp1"], view_type=MetricViewType.TRACES, metric_name="latency", aggregations=[MetricAggregation(aggregation_type=AggregationType.AVG)], dimensions=None, filters=None, time_interval_seconds=None, start_time_ms=None, end_time_ms=None, max_results=MAX_RESULTS_QUERY_TRACE_METRICS, page_token=None, ) assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data == {} def test_invoke_scorer_missing_experiment_id(): with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/scorer/invoke", json={"serialized_scorer": "test", "trace_ids": ["trace1"]}, ) assert response.status_code == 400 data = response.get_json() assert "experiment_id" in data["message"] def test_invoke_scorer_missing_serialized_scorer(): with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/scorer/invoke", json={"experiment_id": "123", "trace_ids": ["trace1"]}, ) assert response.status_code == 400 data = response.get_json() assert "serialized_scorer" in data["message"] def test_invoke_scorer_missing_trace_ids(): with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/scorer/invoke", json={"experiment_id": "123", "serialized_scorer": "test"}, ) assert response.status_code == 400 data = response.get_json() assert "Please select at least one trace to evaluate" in data["message"] def test_invoke_scorer_submits_jobs(mock_tracking_store): serialized_scorer = json.dumps({ "name": "test_judge", "aggregations": [], "description": None, "is_session_level_scorer": False, "mlflow_version": mlflow.__version__, "serialization_version": 1, "builtin_scorer_class": None, "builtin_scorer_pydantic_data": None, "call_source": None, "call_signature": None, "original_func_name": None, "instructions_judge_pydantic_data": { "instructions": "Test: {{ inputs }}", "model": "openai:/gpt-4", "feedback_value_type": { "enum": ["Yes", "No"], "title": "Result", "type": "string", }, }, }) with mock.patch("mlflow.server.jobs.submit_job") as mock_submit: mock_job = mock.MagicMock() mock_job.job_id = "test-job-123" mock_submit.return_value = mock_job with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/scorer/invoke", json={ "experiment_id": "exp-123", "serialized_scorer": serialized_scorer, "trace_ids": ["trace1", "trace2"], }, ) assert response.status_code == 200 data = response.get_json() assert "jobs" in data assert len(data["jobs"]) == 1 assert data["jobs"][0]["job_id"] == "test-job-123" assert data["jobs"][0]["trace_ids"] == ["trace1", "trace2"] mock_submit.assert_called_once() def test_get_ui_telemetry_handler( test_app_context, mock_telemetry_config_cache, bypass_telemetry_env_check ): config = { "disable_telemetry": False, "disable_ui_telemetry": False, "disable_ui_events": ["event1", "event2"], "ui_rollout_percentage": 50, } with mock.patch( "mlflow.server.handlers.fetch_ui_telemetry_config", return_value=config ) as mock_fetch: response = get_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["disable_ui_telemetry"] is False assert response_data["disable_ui_events"] == ["event1", "event2"] # rollout percent gets converted to a float as that is the proto definition assert response_data["ui_rollout_percentage"] == 50.0 assert "config" in mock_telemetry_config_cache assert mock_fetch.call_count == 1 mock_fetch.reset_mock() # subsequent call should hit cache response = get_ui_telemetry_handler() mock_fetch.assert_not_called() assert response_data["disable_ui_telemetry"] is False assert response_data["disable_ui_events"] == ["event1", "event2"] assert response_data["ui_rollout_percentage"] == 50.0 def test_get_ui_telemetry_handler_disabled_by_config( test_app_context, mock_telemetry_config_cache, bypass_telemetry_env_check ): config = { "disable_telemetry": True, "disable_ui_telemetry": False, "disable_ui_events": [], "ui_rollout_percentage": 0, } with mock.patch( "mlflow.server.handlers.fetch_ui_telemetry_config", return_value=config ) as mock_fetch: response = get_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) # if disable_telemetry is True, the server should always report # that UI telemetry is disabled regardless of disable_ui_telemetry assert response_data["disable_ui_telemetry"] is True assert response_data["ui_rollout_percentage"] == 0.0 assert response_data["disable_ui_events"] == [] assert mock_fetch.call_count == 1 def test_get_ui_telemetry_handler_disabled_by_env( test_app_context, mock_telemetry_config_cache, bypass_telemetry_env_check, monkeypatch ): monkeypatch.setenv("DO_NOT_TRACK", "true") with mock.patch("mlflow.server.handlers.fetch_ui_telemetry_config") as mock_fetch: response = get_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) # if telemetry is disabled by env var, the server should always report # that UI telemetry is disabled, and no config fetch should happen mock_fetch.assert_not_called() assert response_data["disable_ui_telemetry"] is True assert response_data["ui_rollout_percentage"] == 0.0 assert response_data["disable_ui_events"] == [] def test_get_ui_telemetry_handler_fallback_values( test_app_context, mock_telemetry_config_cache, bypass_telemetry_env_check ): config_without_ui_fields = { "disable_telemetry": False, "rollout_percentage": 100, } # test fallback values if we forget to define UI config fields with mock.patch("requests.get", return_value=config_without_ui_fields): response = get_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["disable_ui_telemetry"] is True assert response_data["ui_rollout_percentage"] == 0 assert response_data["disable_ui_events"] == [] # test fallback values if we fail to fetch the config with mock.patch("requests.get", return_value=mock.Mock(status_code=404)): response = get_ui_telemetry_handler() assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["disable_ui_telemetry"] is True assert response_data["ui_rollout_percentage"] == 0 assert response_data["disable_ui_events"] == [] def test_post_ui_telemetry_handler_success( test_app, mock_telemetry_config_cache, bypass_telemetry_env_check ): event1 = { "event_name": "test_event_1", "timestamp_ns": 1234567890000000, "params": {"key1": "value1"}, "installation_id": "install-123", "session_id": "session-456", } event2 = { "event_name": "test_event_2", "timestamp_ns": 1234567890000001, "params": {"key2": "value2"}, "installation_id": "install-123", "session_id": "session-456", } request = json.dumps({"records": [event1, event2]}) config = {"disable_ui_telemetry": False, "disable_telemetry": False} mock_client = mock.MagicMock() server_install_id = "server-install-789" with ( test_app.test_request_context( "/ui-telemetry", method="POST", data=request, content_type="application/json" ), mock.patch("mlflow.server.handlers.fetch_ui_telemetry_config", return_value=config), mock.patch("mlflow.server.handlers.get_telemetry_client", return_value=mock_client), mock.patch( "mlflow.server.handlers.get_or_create_installation_id", return_value=server_install_id, ), ): response = post_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["status"] == "success" assert mock_client.add_records.call_count == 1 assert mock_client.add_records.call_args[0][0] == [ Record( **event1, duration_ms=0, status=Status.SUCCESS, server_installation_id=server_install_id, ), Record( **event2, duration_ms=0, status=Status.SUCCESS, server_installation_id=server_install_id, ), ] def test_post_ui_telemetry_handler_telemetry_disabled_by_config( test_app, mock_telemetry_config_cache, bypass_telemetry_env_check ): event = { "event_name": "test_event_1", "timestamp_ns": 1234567890000000, "params": {"key1": "value1"}, "installation_id": "install-123", "session_id": "session-456", } request = json.dumps({"records": [event]}) config = {"disable_ui_telemetry": True} mock_client = mock.MagicMock() with ( test_app.test_request_context( "/ui-telemetry", method="POST", data=request, content_type="application/json" ), mock.patch("mlflow.server.handlers.fetch_ui_telemetry_config", return_value=config), mock.patch("mlflow.server.handlers.get_telemetry_client", return_value=mock_client), ): response = post_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["status"] == "disabled" mock_client.add_record.assert_not_called() def test_post_ui_telemetry_handler_telemetry_disabled_by_env( test_app, mock_telemetry_config_cache, bypass_telemetry_env_check, monkeypatch ): monkeypatch.setenv("DO_NOT_TRACK", "true") request = json.dumps({"records": []}) with ( test_app.test_request_context( "/ui-telemetry", method="POST", data=request, content_type="application/json" ), mock.patch("mlflow.server.handlers.fetch_ui_telemetry_config") as mock_fetch, mock.patch("mlflow.server.handlers.get_telemetry_client") as mock_get_client, ): response = post_ui_telemetry_handler() assert response is not None assert response.status_code == 200 response_data = json.loads(response.get_data()) assert response_data["status"] == "disabled" # assert that no fetch happens and no client is retrieved mock_fetch.assert_not_called() mock_get_client.assert_not_called() def test_send_artifact_prefers_local_path(tmp_path): artifact_path = "test_model/model.pkl" test_data = b"local artifact" test_file = tmp_path / "model.pkl" test_file.write_bytes(test_data) mock_artifact_repo = mock.MagicMock() mock_artifact_repo.get_local_path.return_value = str(test_file) with ( app.test_request_context(method="GET"), mock.patch("mlflow.server.handlers.tempfile.TemporaryDirectory") as mock_tmp_dir, ): response = _send_artifact(mock_artifact_repo, artifact_path) response.direct_passthrough = False assert response.get_data() == test_data assert response.headers["Content-Disposition"] == "attachment; filename=model.pkl" mock_artifact_repo.get_local_path.assert_called_once_with(artifact_path) mock_artifact_repo.download_artifacts.assert_not_called() mock_tmp_dir.assert_not_called() def test_send_artifact_falls_back_to_download_when_local_path_unavailable(tmp_path): artifact_path = "test_model/model.pkl" test_data = b"downloaded artifact" test_file = tmp_path / "model.pkl" test_file.write_bytes(test_data) mock_artifact_repo = mock.MagicMock() mock_artifact_repo.get_local_path.return_value = None mock_artifact_repo.download_artifacts.return_value = str(test_file) with app.test_request_context(method="GET"): response = _send_artifact(mock_artifact_repo, artifact_path) response.direct_passthrough = False assert response.get_data() == test_data assert response.headers["Content-Disposition"] == "attachment; filename=model.pkl" mock_artifact_repo.get_local_path.assert_called_once_with(artifact_path) mock_artifact_repo.download_artifacts.assert_called_once_with(artifact_path, dst_path=mock.ANY) def test_create_artifact_file_response_uses_local_path_mimetype_and_artifact_name(tmp_path): test_file = tmp_path / "payload.html" test_file.write_text("ok") with app.test_request_context(method="GET"): response = _create_artifact_file_response(str(test_file), "artifacts/model.txt") assert response.mimetype == "text/html" assert response.headers["Content-Disposition"] == "attachment; filename=model.txt" def test_create_artifact_file_response_quotes_token_unsafe_ascii_artifact_name(tmp_path): test_file = tmp_path / "payload.html" test_file.write_text("ok") with app.test_request_context(method="GET"): response = _create_artifact_file_response(str(test_file), "artifacts/my model;a.txt") assert response.headers["Content-Disposition"] == 'attachment; filename="my model;a.txt"' def test_download_artifact_uses_local_path_fast_path(enable_serve_artifacts, tmp_path): artifact_path = "test_model/model.pkl" test_data = b"local artifact" test_file = tmp_path / "model.pkl" test_file.write_bytes(test_data) with ( app.test_request_context(method="GET"), mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers.tempfile.TemporaryDirectory") as mock_tmp_dir, ): mock_tmp_dir_instance = mock.MagicMock() mock_tmp_dir.return_value = mock_tmp_dir_instance mock_artifact_repo = mock.MagicMock() mock_artifact_repo.get_local_path.return_value = str(test_file) mock_repo.return_value = mock_artifact_repo response = _download_artifact(artifact_path) response.direct_passthrough = False assert response.get_data() == test_data assert response.headers["Content-Disposition"] == "attachment; filename=model.pkl" mock_artifact_repo.get_local_path.assert_called_once_with(artifact_path) mock_artifact_repo.download_artifacts.assert_not_called() mock_tmp_dir.assert_not_called() def test_upload_artifact_uses_stream_upload_when_mixin_supported(enable_serve_artifacts): artifact_path = "nested/model.pkl" test_data = b"streamed artifact" with ( app.test_request_context( method="PUT", data=test_data, content_type="application/octet-stream" ), mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, ): mock_artifact_repo = mock.MagicMock(spec=LocalArtifactRepository) mock_repo.return_value = mock_artifact_repo response = _upload_artifact(artifact_path) mock_artifact_repo.log_artifact_from_stream.assert_called_once() args, kwargs = mock_artifact_repo.log_artifact_from_stream.call_args assert args[1] == "model.pkl" assert kwargs["artifact_path"] == "nested" assert response.status_code == 200 def test_upload_artifact_falls_back_to_log_artifact_without_mixin(enable_serve_artifacts): artifact_path = "nested/model.pkl" test_data = b"streamed artifact" with ( app.test_request_context( method="PUT", data=test_data, content_type="application/octet-stream" ), mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, ): mock_artifact_repo = mock.MagicMock(spec=ArtifactRepository) mock_repo.return_value = mock_artifact_repo response = _upload_artifact(artifact_path) mock_artifact_repo.log_artifact.assert_called_once() args, kwargs = mock_artifact_repo.log_artifact.call_args assert args[0].endswith("model.pkl") assert kwargs["artifact_path"] == "nested" assert response.status_code == 200 def test_download_artifact_streams_in_chunks(enable_serve_artifacts, tmp_path): # Create a test file with binary data larger than the chunk size (2MB + 1000 bytes) test_file_size = ARTIFACT_STREAM_CHUNK_SIZE * 2 + 1000 test_data = b"x" * test_file_size artifact_path = "test_model/model.pkl" test_file = tmp_path / "model.pkl" test_file.write_bytes(test_data) with ( app.test_request_context(method="GET"), mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers.tempfile.TemporaryDirectory") as mock_tmp_dir, ): # Setup mocks mock_tmp_dir_instance = mock.MagicMock() mock_tmp_dir_instance.name = str(tmp_path) mock_tmp_dir.return_value = mock_tmp_dir_instance mock_artifact_repo = mock.MagicMock() mock_artifact_repo.get_local_path.return_value = None mock_artifact_repo.download_artifacts.return_value = str(test_file) mock_repo.return_value = mock_artifact_repo # Call the function and capture the response response = _download_artifact(artifact_path) # Extract chunks from the response by iterating over its data response_chunks = list(response.response) # Verify that data was streamed in chunks, not line by line # For a 2MB+ binary file, line-by-line would produce many small chunks # Chunk-based streaming should produce exactly 3 chunks (2*1MB + 1000 bytes) assert len(response_chunks) == 3, f"Expected 3 chunks, got {len(response_chunks)}" # Verify chunk sizes assert len(response_chunks[0]) == ARTIFACT_STREAM_CHUNK_SIZE assert len(response_chunks[1]) == ARTIFACT_STREAM_CHUNK_SIZE assert len(response_chunks[2]) == 1000 # Verify that all data is correctly streamed streamed_data = b"".join(response_chunks) assert streamed_data == test_data mock_artifact_repo.get_local_path.assert_called_once_with(artifact_path) mock_artifact_repo.download_artifacts.assert_called_once_with(artifact_path, str(tmp_path)) def test_download_artifact_cleans_up_tmp_dir_when_download_fails(enable_serve_artifacts, tmp_path): artifact_path = "test_model/model.pkl" with ( app.test_request_context(method="GET"), mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers.tempfile.TemporaryDirectory") as mock_tmp_dir, ): mock_tmp_dir_instance = mock.MagicMock() mock_tmp_dir_instance.name = str(tmp_path) mock_tmp_dir.return_value = mock_tmp_dir_instance mock_artifact_repo = mock.MagicMock() mock_artifact_repo.get_local_path.return_value = None mock_artifact_repo.download_artifacts.side_effect = RuntimeError("boom") mock_repo.return_value = mock_artifact_repo with pytest.raises(RuntimeError, match="boom"): _download_artifact(artifact_path) mock_artifact_repo.get_local_path.assert_called_once_with(artifact_path) mock_artifact_repo.download_artifacts.assert_called_once_with(artifact_path, str(tmp_path)) mock_tmp_dir_instance.cleanup.assert_called_once() def test_download_artifact_returns_404_for_missing_azure_blob(enable_serve_artifacts): ResourceNotFoundError = pytest.importorskip("azure.core.exceptions").ResourceNotFoundError artifact_repo = AzureBlobArtifactRepository( "wasbs://container@account.blob.core.windows.net/root", client=mock.MagicMock(), ) artifact_repo.client.get_container_client().walk_blobs.return_value = [] artifact_repo.client.get_container_client().download_blob.side_effect = ResourceNotFoundError( "Operation returned an invalid status: BlobNotFound" ) with ( app.test_request_context(method="GET"), mock.patch( "mlflow.server.handlers._get_artifact_repo_mlflow_artifacts", return_value=artifact_repo, ), ): response = _download_artifact("missing.txt") artifact_repo.thread_pool.shutdown() assert response.status_code == 404 assert json.loads(response.get_data())["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) @pytest.mark.parametrize( ("file_path", "expected_simple", "expected_quoted"), [ # No-extension fully-CJK filename: NFKD normalization strips every # character, so the helper must fall back to a safe non-empty # ``filename=`` value rather than emitting ``filename=;``. ("日本語", "download", "%E6%97%A5%E6%9C%AC%E8%AA%9E"), ("Tribeč_mountains.html", "Tribec_mountains.html", "Tribe%C4%8D_mountains.html"), ( "time_series_eeeúaaa_aaaaaal_39.html", "time_series_eeeuaaa_aaaaaal_39.html", "time_series_eee%C3%BAaaa_aaaaaal_39.html", ), ("日本語.txt", ".txt", "%E6%97%A5%E6%9C%AC%E8%AA%9E.txt"), ], ) def test_response_with_file_attachment_headers_encodes_non_ascii_filename( file_path, expected_simple, expected_quoted ): with app.test_request_context(): response = _response_with_file_attachment_headers(file_path, Response()) header = response.headers["Content-Disposition"] # The Content-Disposition header value must be ASCII-encodable so that # WSGI/ASGI adapters (e.g., starlette's WSGIMiddleware) can serialize # the response without raising UnicodeEncodeError. See GH-23208. header.encode("ascii") assert header == f"attachment; filename={expected_simple}; filename*=UTF-8''{expected_quoted}" @pytest.mark.parametrize( ("filename", "expected_header"), [ ("model.pkl", "attachment; filename=model.pkl"), ("my model;a.txt", 'attachment; filename="my model;a.txt"'), ], ) def test_response_with_file_attachment_headers_ascii_filename_preserves_werkzeug_quoting( filename, expected_header ): with app.test_request_context(): response = _response_with_file_attachment_headers(filename, Response()) assert response.headers["Content-Disposition"] == expected_header @pytest.mark.parametrize( ("filename", "expected_simple", "expected_quoted"), [ # See sibling unit test for why this empty-fallback case leads. ("日本語", "download", "%E6%97%A5%E6%9C%AC%E8%AA%9E"), ("Tribeč_mountains.html", "Tribec_mountains.html", "Tribe%C4%8D_mountains.html"), ( "time_series_eeeúaaa_aaaaaal_39.html", "time_series_eeeuaaa_aaaaaal_39.html", "time_series_eee%C3%BAaaa_aaaaaal_39.html", ), ("日本語.txt", ".txt", "%E6%97%A5%E6%9C%AC%E8%AA%9E.txt"), ], ) def test_download_artifact_endpoint_non_ascii_filename( enable_serve_artifacts, monkeypatch, tmp_path, filename, expected_simple, expected_quoted ): # End-to-end coverage for the `_download_artifact` HTTP path. Exercises the # full `/api/2.0/mlflow-artifacts/artifacts/` route to guard # against regressions where the WSGI/ASGI adapter serializes the # `Content-Disposition` header and rejects raw non-ASCII bytes. See GH-23208. artifact_root = tmp_path / "artifacts" artifact_root.mkdir() file_contents = b"hello from " + filename.encode("utf-8") (artifact_root / filename).write_bytes(file_contents) # ``.as_uri()`` not ``str()``: on Windows ``str(WindowsPath)`` is ``C:\...`` # which mlflow's artifact registry parses as scheme=``C`` and 500s. monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, artifact_root.as_uri()) monkeypatch.setattr("mlflow.server.handlers._artifact_repo", None) quoted_path = urllib.parse.quote(filename) with app.test_client() as c: response = c.get(f"/api/2.0/mlflow-artifacts/artifacts/{quoted_path}") assert response.status_code == 200 assert response.get_data() == file_contents header = response.headers["Content-Disposition"] header.encode("ascii") assert header == f"attachment; filename={expected_simple}; filename*=UTF-8''{expected_quoted}" def test_create_prompt_optimization_job(mock_tracking_store): mock_job_entity = JobEntity( job_id="job-123", creation_time=1234567890, job_name="optimize_prompts", params='{"run_id": "run-456"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890, status_details=None, ) mock_run = mock.MagicMock() mock_run.info.run_id = "run-456" mock_tracking_store.create_run.return_value = mock_run mock_dataset = mock.MagicMock() mock_dataset._to_mlflow_entity.return_value = mock.MagicMock() with ( mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job_entity), mock.patch("mlflow.server.handlers._get_user", return_value="test_user"), mock.patch( "mlflow.genai.datasets.get_dataset", return_value=mock_dataset ) as mock_get_dataset, ): with app.test_request_context( method="POST", json={ "experiment_id": "exp-123", "source_prompt_uri": "prompts:/my-prompt/1", "config": { "optimizer_type": OPTIMIZER_TYPE_GEPA, "dataset_id": "dataset-123", "scorers": ["Correctness", "Safety"], "optimizer_config_json": '{"reflection_model": "openai:/gpt-4"}', }, "tags": [{"key": "env", "value": "test"}], }, ): response = _create_prompt_optimization_job() mock_get_dataset.assert_called_once_with(dataset_id="dataset-123") mock_tracking_store.create_run.assert_called_once() call_kwargs = mock_tracking_store.create_run.call_args[1] assert call_kwargs["experiment_id"] == "exp-123" assert call_kwargs["user_id"] == "test_user" mock_tracking_store.log_batch.assert_called_once() logged_params = mock_tracking_store.log_batch.call_args[1]["params"] param_dict = {p.key: p.value for p in logged_params} assert param_dict["source_prompt_uri"] == "prompts:/my-prompt/1" assert param_dict["optimizer_type"] == "gepa" assert param_dict["dataset_id"] == "dataset-123" assert param_dict["scorer_names"] == '["Correctness", "Safety"]' response_data = json.loads(response.get_data()) assert response_data["job"]["job_id"] == "job-123" assert response_data["job"]["run_id"] == "run-456" assert response_data["job"]["state"]["status"] == "JOB_STATUS_PENDING" assert response_data["job"]["experiment_id"] == "exp-123" assert response_data["job"]["source_prompt_uri"] == "prompts:/my-prompt/1" def test_create_prompt_optimization_job_zero_shot(mock_tracking_store): mock_job_entity = JobEntity( job_id="job-999", creation_time=1234567890, job_name="optimize_prompts", params='{"run_id": "run-999"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890, status_details=None, ) mock_run = mock.MagicMock() mock_run.info.run_id = "run-999" mock_tracking_store.create_run.return_value = mock_run with ( mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job_entity), mock.patch("mlflow.server.handlers._get_user", return_value="test_user"), ): with app.test_request_context( method="POST", json={ "experiment_id": "exp-123", "source_prompt_uri": "prompts:/my-prompt/1", "config": { "optimizer_type": OPTIMIZER_TYPE_METAPROMPT, "scorers": [], # Empty scorers for zero-shot # No dataset_id - zero-shot optimization }, }, ): response = _create_prompt_optimization_job() response_data = json.loads(response.get_data()) assert response_data["job"]["job_id"] == "job-999" assert response_data["job"]["run_id"] == "run-999" assert response_data["job"]["state"]["status"] == "JOB_STATUS_PENDING" mock_tracking_store.log_batch.assert_called_once() logged_params = mock_tracking_store.log_batch.call_args[1]["params"] param_dict = {p.key: p.value for p in logged_params} assert param_dict["dataset_id"] == "" # Empty string for None assert param_dict["scorer_names"] == "[]" # Empty list def test_create_prompt_optimization_job_missing_prompt_uri(mock_tracking_store): with app.test_request_context( method="POST", json={ "experiment_id": "exp-123", "config": { "optimizer_type": 1, "dataset_id": "dataset-123", "scorers": ["Correctness"], }, }, ): response = _create_prompt_optimization_job() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert "source_prompt_uri" in json_response["message"] def test_create_prompt_optimization_job_unspecified_optimizer_type(mock_tracking_store): with app.test_request_context( method="POST", json={ "experiment_id": "exp-123", "source_prompt_uri": "prompts:/my-prompt/1", "config": { "optimizer_type": OPTIMIZER_TYPE_UNSPECIFIED, "dataset_id": "dataset-123", "scorers": ["Correctness"], }, }, ): response = _create_prompt_optimization_job() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert "optimizer_type is required" in json_response["message"] def test_create_prompt_optimization_job_invalid_optimizer_config_json(mock_tracking_store): with app.test_request_context( method="POST", json={ "experiment_id": "exp-123", "source_prompt_uri": "prompts:/my-prompt/1", "config": { "optimizer_type": 1, "dataset_id": "dataset-123", "scorers": ["Correctness"], "optimizer_config_json": "invalid json {", }, }, ): response = _create_prompt_optimization_job() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert "Invalid JSON in optimizer_config_json" in json_response["message"] def test_create_prompt_optimization_job_missing_experiment_id(mock_tracking_store): with app.test_request_context( method="POST", json={ "experiment_id": "", # Empty experiment_id "source_prompt_uri": "prompts:/my-prompt/1", "config": { "optimizer_type": 1, "dataset_id": "dataset-123", "scorers": ["Correctness"], }, }, ): response = _create_prompt_optimization_job() assert response.status_code == 400 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(INVALID_PARAMETER_VALUE) assert "experiment_id is required" in json_response["message"] def test_cancel_prompt_optimization_job(): mock_job_entity = JobEntity( job_id="job-123", creation_time=1234567890, job_name="optimize_prompts", params=( '{"experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1", ' '"run_id": "run-456"}' ), timeout=None, status=JobStatus.CANCELED, result=None, retry_count=0, last_update_time=1234567890, status_details=None, ) with ( mock.patch("mlflow.server.jobs.cancel_job", return_value=mock_job_entity), mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, ): mock_tracking_store = mock.Mock() mock_store.return_value = mock_tracking_store with app.test_request_context(method="POST"): response = _cancel_prompt_optimization_job("job-123") # Verify that the underlying run was terminated mock_tracking_store.update_run_info.assert_called_once() call_args = mock_tracking_store.update_run_info.call_args assert call_args.kwargs["run_id"] == "run-456" assert call_args.kwargs["run_status"] == RunStatus.KILLED assert call_args.kwargs["run_name"] is None assert "end_time" in call_args.kwargs response_data = json.loads(response.get_data()) assert response_data["job"]["job_id"] == "job-123" assert response_data["job"]["state"]["status"] == "JOB_STATUS_CANCELED" assert response_data["job"]["experiment_id"] == "exp-123" assert response_data["job"]["source_prompt_uri"] == "prompts:/my-prompt/1" assert response_data["job"]["run_id"] == "run-456" def test_get_prompt_optimization_job_pending(mock_tracking_store): mock_job = _create_mock_job(status_name="PENDING") mock_run = _create_mock_run( params={ "source_prompt_uri": "prompts:/my-prompt/1", "optimizer_type": "gepa", "dataset_id": "dataset-789", "scorer_names": '["Correctness"]', } ) mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() assert "job" in data job = data["job"] assert job["job_id"] == "job-123" assert job["run_id"] == "run-456" assert job["experiment_id"] == "exp-123" assert job["source_prompt_uri"] == "prompts:/my-prompt/1" assert job["state"]["status"] == "JOB_STATUS_PENDING" def test_get_prompt_optimization_job_succeeded_with_result(mock_tracking_store): mock_job = _create_mock_job( status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/my-prompt/2"}, ) mock_run = _create_mock_run( params={ "source_prompt_uri": "prompts:/my-prompt/1", "optimizer_type": "gepa", "dataset_id": "dataset-789", "scorer_names": '["Correctness", "Safety"]', }, metrics={ "initial_eval_score.Correctness": 0.65, "initial_eval_score.Safety": 0.80, "final_eval_score.Correctness": 0.89, "final_eval_score.Safety": 0.95, }, ) mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] assert job["state"]["status"] == "JOB_STATUS_COMPLETED" assert job["optimized_prompt_uri"] == "prompts:/my-prompt/2" # Verify metrics are populated from the run assert job["initial_eval_scores"]["Correctness"] == 0.65 assert job["initial_eval_scores"]["Safety"] == 0.80 assert job["final_eval_scores"]["Correctness"] == 0.89 assert job["final_eval_scores"]["Safety"] == 0.95 def test_get_prompt_optimization_job_succeeded_run_fetch_fails(mock_tracking_store): mock_job = _create_mock_job( status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/my-prompt/2"}, ) # Simulate run fetch failing (e.g., run not found) mock_tracking_store.get_run.side_effect = Exception("Run not found") with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] assert job["state"]["status"] == "JOB_STATUS_COMPLETED" assert job["optimized_prompt_uri"] == "prompts:/my-prompt/2" # Metrics should not be present when run fetch fails assert "initial_eval_scores" not in job or job["initial_eval_scores"] == {} def test_get_prompt_optimization_job_failed_with_error(mock_tracking_store): mock_job = _create_mock_job( status_name="FAILED", result="Optimization failed: Invalid scorer", ) mock_run = _create_mock_run() mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] assert job["state"]["status"] == "JOB_STATUS_FAILED" assert "Optimization failed" in job["state"]["error_message"] def test_get_prompt_optimization_job_without_run_id(mock_tracking_store): mock_job = _create_mock_job( params={"experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1"} ) with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] assert job["job_id"] == "job-123" assert job["experiment_id"] == "exp-123" assert "run_id" not in job # run_id is not set def test_get_prompt_optimization_job_with_progress(mock_tracking_store): mock_job = _create_mock_job( status_name="RUNNING", params={ "experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1", "run_id": "run-456", "optimizer_config": {"max_metric_calls": 200, "reflection_model": "openai:/gpt-4o"}, }, ) mock_run = _create_mock_run( params={ "source_prompt_uri": "prompts:/my-prompt/1", "optimizer_type": "gepa", }, metrics={ "total_metric_calls": 86, "eval_score": 0.75, }, ) mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] assert job["state"]["status"] == "JOB_STATUS_IN_PROGRESS" # Progress should be 86 / 200 = 0.43 assert job["state"]["metadata"]["progress"] == "0.43" def test_get_prompt_optimization_job_progress_capped_at_one(mock_tracking_store): mock_job = _create_mock_job( status_name="RUNNING", params={ "experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1", "run_id": "run-456", "optimizer_config": {"max_metric_calls": 100, "reflection_model": "openai:/gpt-4o"}, }, ) mock_run = _create_mock_run( metrics={ "total_metric_calls": 150, # Exceeds max_metric_calls }, ) mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] # Progress should be capped at 1.0, not 1.5 assert job["state"]["metadata"]["progress"] == "1.0" def test_get_prompt_optimization_job_no_progress_without_max_metric_calls(mock_tracking_store): mock_job = _create_mock_job( status_name="RUNNING", params={ "experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1", "run_id": "run-456", "optimizer_config": {"reflection_model": "openai:/gpt-4o"}, }, ) mock_run = _create_mock_run( metrics={ "total_metric_calls": 50, }, ) mock_tracking_store.get_run.return_value = mock_run with mock.patch("mlflow.server.jobs.get_job", return_value=mock_job): with app.test_client() as c: response = c.get("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 data = response.get_json() job = data["job"] # Progress should NOT be set when max_metric_calls is not configured assert "progress" not in job["state"].get("status_details", {}) def test_search_prompt_optimization_jobs_returns_multiple_jobs(mock_job_store): mock_jobs = [ _create_mock_job( job_id="job-1", status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/opt/1"}, ), _create_mock_job(job_id="job-2", status_name="RUNNING"), _create_mock_job(job_id="job-3", status_name="PENDING"), ] mock_job_store.list_jobs.return_value = iter(mock_jobs) with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/prompt-optimization/jobs/search", json={"experiment_id": "exp-123"}, ) assert response.status_code == 200 data = response.get_json() assert "jobs" in data assert len(data["jobs"]) == 3 job_ids = [job["job_id"] for job in data["jobs"]] assert "job-1" in job_ids assert "job-2" in job_ids assert "job-3" in job_ids mock_job_store.list_jobs.assert_called_once_with( job_name="optimize_prompts", params={"experiment_id": "exp-123"}, ) def test_search_prompt_optimization_jobs_returns_empty_list(mock_job_store): mock_job_store.list_jobs.return_value = iter([]) with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/prompt-optimization/jobs/search", json={"experiment_id": "exp-456"}, ) assert response.status_code == 200 data = response.get_json() assert data.get("jobs", []) == [] def test_search_prompt_optimization_jobs_missing_experiment_id(): with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/prompt-optimization/jobs/search", json={}, ) assert response.status_code == 400 def test_search_prompt_optimization_jobs_includes_succeeded_job_result(mock_job_store): mock_job = _create_mock_job( job_id="job-1", status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/optimized/1"}, ) mock_job_store.list_jobs.return_value = iter([mock_job]) with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/prompt-optimization/jobs/search", json={"experiment_id": "exp-123"}, ) assert response.status_code == 200 data = response.get_json() assert len(data["jobs"]) == 1 assert data["jobs"][0]["optimized_prompt_uri"] == "prompts:/optimized/1" def test_search_prompt_optimization_jobs_includes_failed_job_error(mock_job_store): mock_job = _create_mock_job( job_id="job-1", status_name="FAILED", result="Some error occurred", ) mock_job_store.list_jobs.return_value = iter([mock_job]) with app.test_client() as c: response = c.post( "/ajax-api/3.0/mlflow/prompt-optimization/jobs/search", json={"experiment_id": "exp-123"}, ) assert response.status_code == 200 data = response.get_json() assert len(data["jobs"]) == 1 assert "Some error occurred" in data["jobs"][0]["state"]["error_message"] def test_delete_prompt_optimization_job_success(mock_job_store, mock_tracking_store): mock_job = _create_mock_job( status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/optimized/1"}, ) mock_job_store.get_job.return_value = mock_job mock_tracking_store.get_run.return_value = mock.MagicMock() # Run exists with app.test_client() as c: response = c.delete("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 mock_job_store.delete_jobs.assert_called_once_with(job_ids=["job-123"]) mock_tracking_store.get_run.assert_called_once_with("run-456") mock_tracking_store.delete_run.assert_called_once_with("run-456") def test_delete_prompt_optimization_job_without_run_id(mock_job_store, mock_tracking_store): mock_job = _create_mock_job( params={"experiment_id": "exp-123", "prompt_uri": "prompts:/my-prompt/1"} ) mock_job_store.get_job.return_value = mock_job with app.test_client() as c: response = c.delete("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 mock_job_store.delete_jobs.assert_called_once_with(job_ids=["job-123"]) mock_tracking_store.delete_run.assert_not_called() def test_delete_prompt_optimization_job_skips_run_deletion_when_run_not_found( mock_job_store, mock_tracking_store ): mock_job = _create_mock_job( status_name="SUCCEEDED", result={"optimized_prompt_uri": "prompts:/optimized/1"}, ) mock_job_store.get_job.return_value = mock_job mock_tracking_store.get_run.side_effect = MlflowException("Run not found") with app.test_client() as c: response = c.delete("/ajax-api/3.0/mlflow/prompt-optimization/jobs/job-123") assert response.status_code == 200 mock_job_store.delete_jobs.assert_called_once_with(job_ids=["job-123"]) # delete_run should not be called since run doesn't exist mock_tracking_store.delete_run.assert_not_called() def test_get_workspace_scoped_repo_path_if_enabled_allows_legacy_default_artifacts(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext(DEFAULT_WORKSPACE_NAME): assert ( _get_workspace_scoped_repo_path_if_enabled("1/legacy/artifact") == "1/legacy/artifact" ) def test_get_workspace_scoped_repo_path_if_enabled_still_scopes_non_default(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext("team-blue"): scoped = _get_workspace_scoped_repo_path_if_enabled("2/new/artifact") assert scoped.startswith("workspaces/team-blue/2/new/artifact") def test_get_workspace_scoped_repo_path_if_enabled_prevents_cross_workspace_access(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext("team-a"): with pytest.raises(MlflowException, match="targets workspace 'team-b'"): _get_workspace_scoped_repo_path_if_enabled("workspaces/team-b/secret.txt") with pytest.raises(MlflowException, match="targets workspace 'other'"): _get_workspace_scoped_repo_path_if_enabled("workspaces/other/data/model.pkl") def test_get_workspace_scoped_repo_path_if_enabled_rejects_empty_workspace_in_path(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext("team-a"): with pytest.raises(MlflowException, match="must include a workspace name"): _get_workspace_scoped_repo_path_if_enabled("workspaces/") with pytest.raises(MlflowException, match="must include a workspace name"): _get_workspace_scoped_repo_path_if_enabled("workspaces//data.txt") def test_get_workspace_scoped_repo_path_if_enabled_allows_matching_workspace_prefix(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext("team-a"): result = _get_workspace_scoped_repo_path_if_enabled("workspaces/team-a/data.txt") assert result == "workspaces/team-a/data.txt" result = _get_workspace_scoped_repo_path_if_enabled("/workspaces/team-a/nested/path") assert result == "workspaces/team-a/nested/path" def test_get_workspace_scoped_repo_path_if_enabled_default_workspace_cross_access_blocked( monkeypatch, ): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with WorkspaceContext(DEFAULT_WORKSPACE_NAME): result = _get_workspace_scoped_repo_path_if_enabled("legacy/artifact.txt") assert result == "legacy/artifact.txt" with pytest.raises(MlflowException, match="targets workspace 'team-b'"): _get_workspace_scoped_repo_path_if_enabled("workspaces/team-b/data.txt") def test_get_workspace_scoped_repo_path_if_enabled_requires_active_workspace(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") with pytest.raises(MlflowException, match="Active workspace is required"): _get_workspace_scoped_repo_path_if_enabled("some/path") def test_get_artifact_handler_applies_workspace_scoping(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_run = mock.MagicMock() mock_run.info.artifact_uri = "mlflow-artifacts:/exp1/run1/artifacts" mock_artifact_repo = mock.MagicMock() mock_artifact_repo.download_artifacts.return_value = "/tmp/artifact.txt" with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers._send_artifact") as mock_send, ): mock_store.return_value.get_run.return_value = mock_run mock_repo.return_value = mock_artifact_repo with WorkspaceContext("team-blue"): with app.test_request_context( method="GET", query_string={"run_id": "run1", "path": "model/weights.bin"} ): get_artifact_handler() mock_send.assert_called_once() artifact_path = mock_send.call_args[0][1] assert artifact_path.startswith("workspaces/team-blue/") def test_get_artifact_handler_no_scoping_when_workspaces_disabled(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "false") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_run = mock.MagicMock() mock_run.info.artifact_uri = "mlflow-artifacts:/exp1/run1/artifacts" mock_artifact_repo = mock.MagicMock() with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers._send_artifact") as mock_send, ): mock_store.return_value.get_run.return_value = mock_run mock_repo.return_value = mock_artifact_repo with app.test_request_context( method="GET", query_string={"run_id": "run1", "path": "model/weights.bin"} ): get_artifact_handler() mock_send.assert_called_once() artifact_path = mock_send.call_args[0][1] assert not artifact_path.startswith("workspaces/") def test_get_model_version_artifact_handler_applies_workspace_scoping(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_artifact_repo = mock.MagicMock() with ( mock.patch("mlflow.server.handlers._get_model_registry_store") as mock_store, mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers._send_artifact") as mock_send, ): mock_store.return_value.get_model_version_download_uri.return_value = ( "mlflow-artifacts:/models/MyModel/1" ) mock_repo.return_value = mock_artifact_repo with WorkspaceContext("team-red"): with app.test_request_context( method="GET", query_string={"name": "MyModel", "version": "1", "path": "model.pkl"} ): get_model_version_artifact_handler() mock_send.assert_called_once() artifact_path = mock_send.call_args[0][1] assert artifact_path.startswith("workspaces/team-red/") def test_get_logged_model_artifact_handler_applies_workspace_scoping(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_logged_model = mock.MagicMock() mock_logged_model.artifact_location = "mlflow-artifacts:/exp1/run1/artifacts/model" mock_artifact_repo = mock.MagicMock() with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, mock.patch("mlflow.server.handlers._send_artifact") as mock_send, ): mock_store.return_value.get_logged_model.return_value = mock_logged_model mock_repo.return_value = mock_artifact_repo with WorkspaceContext("team-green"): with app.test_request_context( method="GET", query_string={"artifact_file_path": "MLmodel"} ): get_logged_model_artifact_handler("model123") mock_send.assert_called_once() artifact_path = mock_send.call_args[0][1] assert artifact_path.startswith("workspaces/team-green/") def test_upload_artifact_handler_applies_workspace_scoping(monkeypatch): monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_run = mock.MagicMock() mock_run.info.artifact_uri = "mlflow-artifacts:/exp1/run1/artifacts" mock_run.info.experiment_id = "exp1" mock_run.info.run_id = "run1" mock_artifact_repo = mock.MagicMock() with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, ): mock_store.return_value.get_run.return_value = mock_run mock_repo.return_value = mock_artifact_repo with WorkspaceContext("team-purple"): with app.test_request_context( method="POST", query_string={"run_uuid": "run1", "path": "output.txt"}, data=b"test data", ): upload_artifact_handler() mock_artifact_repo.log_artifact.assert_called_once() logged_path = mock_artifact_repo.log_artifact.call_args[0][1] assert logged_path.startswith("workspaces/team-purple/") def test_list_artifacts_for_proxied_run_artifact_root_applies_workspace_scoping(monkeypatch): from mlflow.store.artifact.artifact_repo import ArtifactRepository monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true") monkeypatch.setenv(SERVE_ARTIFACTS_ENV_VAR, "true") monkeypatch.setenv(ARTIFACTS_DESTINATION_ENV_VAR, "s3://bucket") mock_artifact_repo = mock.MagicMock(spec=ArtifactRepository) mock_artifact_repo.list_artifacts.return_value = [] with ( mock.patch("mlflow.server.handlers._get_artifact_repo_mlflow_artifacts") as mock_repo, WorkspaceContext("team-orange"), ): mock_repo.return_value = mock_artifact_repo _list_artifacts_for_proxied_run_artifact_root( proxied_artifact_root="mlflow-artifacts:/exp1/run1/artifacts", relative_path="model", ) mock_artifact_repo.list_artifacts.assert_called_once() listed_path = mock_artifact_repo.list_artifacts.call_args[0][0] assert listed_path.startswith("workspaces/team-orange/") # ==================== Budget Window Tests ==================== def _make_budget_policy( budget_policy_id="bp-test", budget_amount=100.0, duration=None, ): return GatewayBudgetPolicy( budget_policy_id=budget_policy_id, budget_unit=BudgetUnit.USD, budget_amount=budget_amount, duration=duration or BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.GLOBAL, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, ) def test_list_budget_windows_empty(): with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker") as mock_tracker, mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), ): mock_tracker.return_value.get_all_windows.return_value = [] response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 assert response.json.get("windows", []) == [] def test_list_budget_windows_returns_window_data(): tracker = InMemoryBudgetTracker() policy = _make_budget_policy(budget_policy_id="bp-1", budget_amount=50.0) tracker.refresh_policies([policy]) tracker.record_cost(12.5) with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker", return_value=tracker), mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), ): response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 data = response.json assert len(data["windows"]) == 1 window = data["windows"][0] assert window["budget_policy_id"] == "bp-1" assert window["current_spend"] == 12.5 min_ms = int(datetime(2000, 1, 1, tzinfo=timezone.utc).timestamp() * 1000) assert window["window_start_ms"] >= min_ms assert window["window_end_ms"] > window["window_start_ms"] # Policy uses duration_unit=DAYS, duration_value=1 → exactly 1 day assert window["window_end_ms"] - window["window_start_ms"] == 86_400_000 def test_list_budget_windows_multiple_policies(): tracker = InMemoryBudgetTracker() policy1 = _make_budget_policy(budget_policy_id="bp-1", budget_amount=100.0) policy2 = _make_budget_policy(budget_policy_id="bp-2", budget_amount=200.0) tracker.refresh_policies([policy1, policy2]) tracker.record_cost(30.0) with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker", return_value=tracker), mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), ): response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 data = response.json policy_ids = {w["budget_policy_id"] for w in data["windows"]} assert policy_ids == {"bp-1", "bp-2"} windows_by_id = {w["budget_policy_id"]: w for w in data["windows"]} assert windows_by_id["bp-1"]["current_spend"] == 30.0 assert windows_by_id["bp-2"]["current_spend"] == 30.0 def test_list_budget_windows_zero_spend(): tracker = InMemoryBudgetTracker() policy = _make_budget_policy(budget_amount=100.0) tracker.refresh_policies([policy]) with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker", return_value=tracker), mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), ): response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 window = response.json["windows"][0] assert window["budget_policy_id"] == "bp-test" assert window["current_spend"] == 0.0 def test_list_budget_windows_workspace_scoped_filters_workspace_policies(): tracker = InMemoryBudgetTracker() global_policy = GatewayBudgetPolicy( budget_policy_id="bp-global", budget_unit=BudgetUnit.USD, budget_amount=100.0, duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.GLOBAL, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, ) ws_policy = GatewayBudgetPolicy( budget_policy_id="bp-ws", budget_unit=BudgetUnit.USD, budget_amount=50.0, duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.WORKSPACE, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, workspace="team-a", ) other_ws_policy = GatewayBudgetPolicy( budget_policy_id="bp-other", budget_unit=BudgetUnit.USD, budget_amount=75.0, duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.WORKSPACE, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, workspace="team-b", ) tracker.refresh_policies([global_policy, ws_policy, other_ws_policy]) with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker", return_value=tracker), mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), WorkspaceContext("team-a"), ): response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 policy_ids = {w["budget_policy_id"] for w in response.json["windows"]} assert policy_ids == {"bp-global", "bp-ws"} def test_list_budget_windows_no_workspace_returns_all(): tracker = InMemoryBudgetTracker() global_policy = GatewayBudgetPolicy( budget_policy_id="bp-global", budget_unit=BudgetUnit.USD, budget_amount=100.0, duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.GLOBAL, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, ) ws_policy = GatewayBudgetPolicy( budget_policy_id="bp-ws", budget_unit=BudgetUnit.USD, budget_amount=50.0, duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1), target_scope=BudgetTargetScope.WORKSPACE, budget_action=BudgetAction.ALERT, created_at=0, last_updated_at=0, workspace="team-a", ) tracker.refresh_policies([global_policy, ws_policy]) with ( app.test_client() as c, mock.patch("mlflow.server.handlers.get_budget_tracker", return_value=tracker), mock.patch("mlflow.server.handlers.maybe_refresh_budget_policies"), ): response = c.get("/ajax-api/3.0/mlflow/gateway/budgets/windows") assert response.status_code == 200 policy_ids = {w["budget_policy_id"] for w in response.json["windows"]} assert policy_ids == {"bp-global", "bp-ws"} def test_create_issue_with_all_fields(): request_message = CreateIssue() request_message.experiment_id = "exp-123" request_message.name = "High latency" request_message.description = "API calls are taking too long" request_message.status = "pending" request_message.source_run_id = "run-123" request_message.root_causes.extend(["Database query inefficiency", "Network latency"]) request_message.categories.extend(["performance", "database"]) request_message.severity = IssueSeverity.HIGH.value request_message.created_by = "user@example.com" issue = Issue( issue_id="iss-123", experiment_id="exp-123", name="High latency", description="API calls are taking too long", status=IssueStatus.PENDING, source_run_id="run-123", root_causes=["Database query inefficiency", "Network latency"], categories=["performance", "database"], severity=IssueSeverity.HIGH, created_timestamp=1234567890, last_updated_timestamp=1234567890, created_by="user@example.com", ) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.create_issue.return_value = issue response = _create_issue() mock_store.return_value.create_issue.assert_called_once() call_kwargs = mock_store.return_value.create_issue.call_args[1] assert call_kwargs["experiment_id"] == "exp-123" assert call_kwargs["name"] == "High latency" assert call_kwargs["description"] == "API calls are taking too long" assert call_kwargs["status"] == IssueStatus.PENDING assert call_kwargs["source_run_id"] == "run-123" assert call_kwargs["root_causes"] == ["Database query inefficiency", "Network latency"] assert call_kwargs["categories"] == ["performance", "database"] assert call_kwargs["severity"] == IssueSeverity.HIGH.value assert call_kwargs["created_by"] == "user@example.com" json_response = json.loads(response.get_data()) assert json_response["issue"]["issue_id"] == "iss-123" assert json_response["issue"]["root_causes"] == [ "Database query inefficiency", "Network latency", ] assert json_response["issue"]["categories"] == ["performance", "database"] def test_create_issue_without_optional_fields(): request_message = CreateIssue() request_message.experiment_id = "exp-456" request_message.name = "Error handling issue" request_message.description = "Errors are not being caught properly" issue = Issue( issue_id="iss-456", experiment_id="exp-456", name="Error handling issue", description="Errors are not being caught properly", status=IssueStatus.PENDING, created_timestamp=1234567890, last_updated_timestamp=1234567890, ) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.create_issue.return_value = issue response = _create_issue() mock_store.return_value.create_issue.assert_called_once() call_kwargs = mock_store.return_value.create_issue.call_args[1] assert call_kwargs["source_run_id"] is None assert call_kwargs["root_causes"] is None assert "severity" not in call_kwargs json_response = json.loads(response.get_data()) assert json_response["issue"]["issue_id"] == "iss-456" def test_create_issue_with_default_status(): request_message = CreateIssue() request_message.experiment_id = "exp-789" request_message.name = "Test issue" request_message.description = "Test description" issue = Issue( issue_id="iss-789", experiment_id="exp-789", name="Test issue", description="Test description", status=IssueStatus.PENDING, created_timestamp=1234567890, last_updated_timestamp=1234567890, ) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.create_issue.return_value = issue _create_issue() call_kwargs = mock_store.return_value.create_issue.call_args[1] # Status should not be in kwargs when not provided (store uses default) assert "status" not in call_kwargs def test_get_issue(): issue = Issue( issue_id="iss-get-123", experiment_id="exp-123", name="Test issue", description="Test description", status=IssueStatus.RESOLVED, severity=IssueSeverity.HIGH, root_causes=["Root cause 1"], created_timestamp=1234567890, last_updated_timestamp=1234567890, ) with mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store: mock_store.return_value.get_issue.return_value = issue with app.test_request_context(): response = _get_issue("iss-get-123") mock_store.return_value.get_issue.assert_called_once_with("iss-get-123") json_response = json.loads(response.get_data()) assert json_response["issue"]["issue_id"] == "iss-get-123" assert json_response["issue"]["name"] == "Test issue" assert json_response["issue"]["severity"] == "high" assert json_response["issue"]["root_causes"] == ["Root cause 1"] def test_get_issue_not_found(): with mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store: mock_store.return_value.get_issue.side_effect = MlflowException( "Issue not found", error_code=RESOURCE_DOES_NOT_EXIST ) with app.test_request_context(): response = _get_issue("nonexistent-id") # The @catch_mlflow_exception decorator catches and returns error as JSON assert response.status_code == 404 json_response = json.loads(response.get_data()) assert json_response["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) assert "Issue not found" in json_response["message"] def test_update_issue(): request_message = UpdateIssue() request_message.issue_id = "iss-update-123" request_message.name = "Updated issue name" request_message.description = "Updated description" request_message.status = "resolved" request_message.severity = "medium" updated_issue = Issue( issue_id="iss-update-123", experiment_id="exp-123", name="Updated issue name", description="Updated description", status=IssueStatus.RESOLVED, severity=IssueSeverity.MEDIUM, created_timestamp=1234567890, last_updated_timestamp=1234567900, ) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.update_issue.return_value = updated_issue response = _update_issue("iss-update-123") mock_store.return_value.update_issue.assert_called_once() call_kwargs = mock_store.return_value.update_issue.call_args[1] assert call_kwargs["issue_id"] == "iss-update-123" assert call_kwargs["name"] == "Updated issue name" assert call_kwargs["description"] == "Updated description" assert call_kwargs["status"] == IssueStatus.RESOLVED assert call_kwargs["severity"] == IssueSeverity.MEDIUM.value json_response = json.loads(response.get_data()) assert json_response["issue"]["issue_id"] == "iss-update-123" assert json_response["issue"]["name"] == "Updated issue name" assert json_response["issue"]["severity"] == "medium" def test_search_issues_all(): request_message = SearchIssues() issues = [ Issue( issue_id="iss-1", experiment_id="exp-1", name="Issue 1", description="Description 1", status=IssueStatus.PENDING, created_timestamp=1234567890, last_updated_timestamp=1234567890, ), Issue( issue_id="iss-2", experiment_id="exp-1", name="Issue 2", description="Description 2", status=IssueStatus.RESOLVED, created_timestamp=1234567891, last_updated_timestamp=1234567891, ), ] with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.search_issues.return_value = PagedList(issues, token="next-token") response = _search_issues() mock_store.return_value.search_issues.assert_called_once() call_kwargs = mock_store.return_value.search_issues.call_args[1] # max_results not specified in request, so it's not passed to store # The store will use its own default parameter value (SEARCH_ISSUES_DEFAULT_MAX_RESULTS) assert "max_results" not in call_kwargs assert call_kwargs["experiment_id"] is None assert call_kwargs["filter_string"] is None json_response = json.loads(response.get_data()) assert len(json_response["issues"]) == 2 assert json_response["issues"][0]["issue_id"] == "iss-1" assert json_response["issues"][1]["issue_id"] == "iss-2" assert json_response["next_page_token"] == "next-token" def test_search_issues_with_filters(): request_message = SearchIssues() request_message.experiment_id = "exp-specific" request_message.filter_string = "status = 'resolved' AND source_run_id = 'run-specific'" request_message.max_results = 50 issues = [ Issue( issue_id="iss-filtered", experiment_id="exp-specific", name="Filtered issue", description="Description", status=IssueStatus.RESOLVED, source_run_id="run-specific", created_timestamp=1234567890, last_updated_timestamp=1234567890, ), ] with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.search_issues.return_value = PagedList(issues, token=None) response = _search_issues() call_kwargs = mock_store.return_value.search_issues.call_args[1] assert call_kwargs["experiment_id"] == "exp-specific" assert ( call_kwargs["filter_string"] == "status = 'resolved' AND source_run_id = 'run-specific'" ) assert call_kwargs["max_results"] == 50 json_response = json.loads(response.get_data()) assert len(json_response["issues"]) == 1 assert json_response["issues"][0]["issue_id"] == "iss-filtered" assert json_response["next_page_token"] == "" def test_search_issues_with_pagination(): request_message = SearchIssues() request_message.max_results = 10 request_message.page_token = "token-123" issues = [ Issue( issue_id=f"iss-{i}", experiment_id="exp-1", name=f"Issue {i}", description=f"Description {i}", status=IssueStatus.PENDING, created_timestamp=1234567890 + i, last_updated_timestamp=1234567890 + i, ) for i in range(10) ] with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.search_issues.return_value = PagedList(issues, token="token-456") response = _search_issues() call_kwargs = mock_store.return_value.search_issues.call_args[1] assert call_kwargs["max_results"] == 10 assert call_kwargs["page_token"] == "token-123" json_response = json.loads(response.get_data()) assert len(json_response["issues"]) == 10 assert json_response["next_page_token"] == "token-456" def test_search_issues_empty_results(): request_message = SearchIssues() with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.search_issues.return_value = PagedList([], token=None) response = _search_issues() json_response = json.loads(response.get_data()) assert len(json_response.get("issues", [])) == 0 assert json_response["next_page_token"] == "" def test_search_issues_with_trace_count(): request_message = SearchIssues() request_message.include_trace_count = True issues = [ Issue( issue_id="iss-1", experiment_id="exp-1", name="Issue with traces", description="Has 2 traces", status=IssueStatus.PENDING, created_timestamp=1234567890, last_updated_timestamp=1234567890, trace_count=2, ), Issue( issue_id="iss-2", experiment_id="exp-1", name="Issue without traces", description="Has no traces", status=IssueStatus.PENDING, created_timestamp=1234567891, last_updated_timestamp=1234567891, trace_count=0, ), ] with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.search_issues.return_value = PagedList(issues, token=None) response = _search_issues() call_kwargs = mock_store.return_value.search_issues.call_args[1] assert call_kwargs["include_trace_count"] is True json_response = json.loads(response.get_data()) assert len(json_response["issues"]) == 2 assert json_response["issues"][0]["trace_count"] == 2 assert json_response["issues"][1]["trace_count"] == 0 def test_create_issue_with_empty_lists(): request_message = CreateIssue() request_message.experiment_id = "exp-123" request_message.name = "Test issue" request_message.description = "Test description" issue = Issue( issue_id="iss-empty-lists", experiment_id="exp-123", name="Test issue", description="Test description", status=IssueStatus.PENDING, created_timestamp=1234567890, last_updated_timestamp=1234567890, ) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.create_issue.return_value = issue _create_issue() call_kwargs = mock_store.return_value.create_issue.call_args[1] # Empty lists should be passed as None assert call_kwargs["root_causes"] is None def test_invoke_issue_detection_handler_success(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_job = JobEntity( job_id="job-123", creation_time=1234567890000, job_name="invoke_issue_detection", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890000, status_details=None, ) mock_run_info = mock.MagicMock() mock_run_info.run_id = "run-123" mock_run = mock.MagicMock() mock_run.info = mock_run_info request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1", "trace-2"], "categories": ["correctness", "safety"], "provider": "openai", "model": "gpt-4o", "secret_id": "secret-123", } with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch( "mlflow.genai.discovery.job._fetch_provider_credentials", return_value={"OPENAI_API_KEY": "test-key"}, ) as mock_fetch_creds, mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job) as mock_submit_job, mock.patch("mlflow.start_run", return_value=mock_run), mock.patch("mlflow.set_tag"), mock.patch("mlflow.end_run"), app.test_client() as c, ): resp = c.post( "/ajax-api/3.0/mlflow/issues/invoke", json=request_json, ) assert resp.status_code == 200 json_response = resp.get_json() assert json_response["job_id"] == "job-123" assert json_response["run_id"] == "run-123" mock_fetch_creds.assert_called_once_with(mock_store.return_value, "openai", "secret-123") mock_submit_job.assert_called_once() call_kwargs = mock_submit_job.call_args.kwargs assert call_kwargs["params"]["experiment_id"] == "exp-123" assert call_kwargs["params"]["trace_ids"] == ["trace-1", "trace-2"] assert call_kwargs["params"]["categories"] == ["correctness", "safety"] assert call_kwargs["params"]["model"] == "openai:/gpt-4o" assert call_kwargs["extra_envs"] == {"OPENAI_API_KEY": "test-key"} def test_invoke_issue_detection_handler_with_endpoint(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_job = JobEntity( job_id="job-456", creation_time=1234567890000, job_name="invoke_issue_detection", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890000, status_details=None, ) mock_run_info = mock.MagicMock() mock_run_info.run_id = "run-456" mock_run = mock.MagicMock() mock_run.info = mock_run_info request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "categories": ["correctness"], "provider": "openai", "endpoint_name": "my-endpoint", "secret_id": "secret-123", } with ( mock.patch( "mlflow.genai.discovery.job._fetch_provider_credentials", return_value={"OPENAI_API_KEY": "test-key"}, ), mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job) as mock_submit_job, mock.patch("mlflow.start_run", return_value=mock_run), mock.patch("mlflow.set_tag"), mock.patch("mlflow.end_run"), app.test_client() as c, ): resp = c.post( "/ajax-api/3.0/mlflow/issues/invoke", json=request_json, ) assert resp.status_code == 200 json_response = resp.get_json() assert json_response["job_id"] == "job-456" assert json_response["run_id"] == "run-456" call_kwargs = mock_submit_job.call_args.kwargs assert call_kwargs["params"]["model"] == "gateway:/my-endpoint" def test_invoke_issue_detection_handler_missing_required_params(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "categories": ["correctness"], "provider": "openai", # Missing both 'model' and 'endpoint_name' "secret_id": "secret-123", } with ( mock.patch( "mlflow.genai.discovery.job._fetch_provider_credentials", return_value={"OPENAI_API_KEY": "test-key"}, ), app.test_client() as c, ): resp = c.post( "/ajax-api/3.0/mlflow/issues/invoke", json=request_json, ) assert resp.status_code == 500 json_response = resp.get_json() assert ( "Either 'endpoint_name' or both 'provider' and 'model' must be provided" in json_response["message"] ) def _make_genai_evaluate_job(job_id: str = "job-genai-1") -> JobEntity: return JobEntity( job_id=job_id, creation_time=1234567890000, job_name="invoke_genai_evaluate", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890000, status_details=None, ) def test_invoke_genai_evaluate_handler_success(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_job = _make_genai_evaluate_job() mock_run = mock.MagicMock() mock_run.info.run_id = "run-genai-1" mock_client = mock.MagicMock() mock_client.create_run.return_value = mock_run request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1", "trace-2"], "serialized_scorers": ['{"name":"my-judge"}', '{"name":"safety"}'], } with ( mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job) as mock_submit_job, mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) assert resp.status_code == 200 json_response = resp.get_json() assert json_response == {"job_id": "job-genai-1", "run_id": "run-genai-1"} # The handler must create the run *upfront* with the right MLFLOW_RUN_TYPE tag # so the run appears on /evaluation-runs immediately, before the job starts # producing artifacts. We use MlflowClient (not mlflow.start_run) so the run # is created directly in the store without touching the fluent active-run stack. mock_client.create_run.assert_called_once_with( experiment_id="exp-123", tags={"mlflow.runType": "genai_evaluate"}, ) mock_submit_job.assert_called_once() submit_kwargs = mock_submit_job.call_args.kwargs assert submit_kwargs["params"]["trace_ids"] == ["trace-1", "trace-2"] assert submit_kwargs["params"]["serialized_scorers"] == request_json["serialized_scorers"] assert submit_kwargs["params"]["run_id"] == "run-genai-1" # No basic auth on the test client -> no username propagated. assert submit_kwargs["params"]["username"] is None mock_client.set_tag.assert_called_once_with( "run-genai-1", "mlflow.genaiEvaluate.jobId", "job-genai-1" ) mock_client.set_terminated.assert_not_called() def test_invoke_genai_evaluate_handler_rejects_empty_trace_ids(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_client = mock.MagicMock() request_json = { "experiment_id": "exp-123", "trace_ids": [], "serialized_scorers": ['{"name":"my-judge"}'], } with ( mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), mock.patch("mlflow.server.jobs.submit_job") as mock_submit_job, app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) assert resp.status_code == 400 assert "Please select at least one trace to evaluate." in resp.get_json()["message"] # Guard against accidentally creating a run for an invalid request. mock_client.create_run.assert_not_called() mock_submit_job.assert_not_called() def test_invoke_genai_evaluate_handler_rejects_empty_serialized_scorers(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_client = mock.MagicMock() request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "serialized_scorers": [], } with ( mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), mock.patch("mlflow.server.jobs.submit_job") as mock_submit_job, app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) assert resp.status_code == 400 assert "Please select at least one judge." in resp.get_json()["message"] mock_client.create_run.assert_not_called() mock_submit_job.assert_not_called() def test_invoke_genai_evaluate_handler_missing_required_fields(monkeypatch): monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") # Each call drops one of the required fields; the schema validator # should reject before we ever create a run. base = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "serialized_scorers": ['{"name":"my-judge"}'], } for missing in ("experiment_id", "trace_ids", "serialized_scorers"): request_json = {k: v for k, v in base.items() if k != missing} mock_client = mock.MagicMock() with ( mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), mock.patch("mlflow.server.jobs.submit_job") as mock_submit_job, app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) assert resp.status_code == 400, f"missing={missing}: {resp.get_json()}" mock_client.create_run.assert_not_called() mock_submit_job.assert_not_called() def test_invoke_genai_evaluate_handler_propagates_basic_auth_username(monkeypatch): """Username comes from HTTP Basic auth and feeds the job's gateway-auth path so judge LLM calls are made *as* the user. """ monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_job = _make_genai_evaluate_job("job-auth") mock_run = mock.MagicMock() mock_run.info.run_id = "run-auth" mock_client = mock.MagicMock() mock_client.create_run.return_value = mock_run request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "serialized_scorers": ['{"name":"my-judge"}'], } with ( mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job) as mock_submit_job, mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), app.test_client() as c, ): # Encoded form of "alice:hunter2" resp = c.post( "/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json, headers={"Authorization": "Basic YWxpY2U6aHVudGVyMg=="}, ) assert resp.status_code == 200 assert mock_submit_job.call_args.kwargs["params"]["username"] == "alice" def test_invoke_genai_evaluate_handler_marks_run_failed_when_submit_job_raises(monkeypatch): """If submit_job raises after the run is created, the handler must flip the run to FAILED itself — otherwise it'd be stuck in RUNNING forever because the worker that would normally do that transition was never enqueued. """ monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_run = mock.MagicMock() mock_run.info.run_id = "run-fail" mock_client = mock.MagicMock() mock_client.create_run.return_value = mock_run request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "serialized_scorers": ['{"name":"my-judge"}'], } submit_error = MlflowException( "Mlflow server job execution feature is not enabled.", error_code=INVALID_PARAMETER_VALUE, ) with ( mock.patch("mlflow.server.jobs.submit_job", side_effect=submit_error), mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) # @catch_mlflow_exception turns the re-raised MlflowException into a 4xx. assert resp.status_code != 200 assert "job execution feature is not enabled" in resp.get_json()["message"] mock_client.set_terminated.assert_called_once_with("run-fail", "FAILED") # The job-id tag must not have been written — the job never existed. mock_client.set_tag.assert_not_called() def test_invoke_genai_evaluate_handler_marks_run_failed_when_set_tag_raises(monkeypatch): """The same try/except must also cover the post-submit set_tag call. If the tag write fails (e.g. transient store error) we'd otherwise leave the run in RUNNING because nothing else writes a terminal status from the handler. """ monkeypatch.setenv("MLFLOW_SERVER_ENABLE_JOB_EXECUTION", "true") mock_job = _make_genai_evaluate_job("job-tag-fail") mock_run = mock.MagicMock() mock_run.info.run_id = "run-tag-fail" mock_client = mock.MagicMock() mock_client.create_run.return_value = mock_run mock_client.set_tag.side_effect = RuntimeError("store unavailable") request_json = { "experiment_id": "exp-123", "trace_ids": ["trace-1"], "serialized_scorers": ['{"name":"my-judge"}'], } with ( mock.patch("mlflow.server.jobs.submit_job", return_value=mock_job), mock.patch("mlflow.server.handlers.MlflowClient", return_value=mock_client), app.test_client() as c, ): resp = c.post("/ajax-api/3.0/mlflow/genai/evaluate/invoke", json=request_json) assert resp.status_code != 200 mock_client.set_terminated.assert_called_once_with("run-tag-fail", "FAILED") def test_get_job_success(mock_job_store): mock_job = JobEntity( job_id="job-123", creation_time=1234567890000, job_name="invoke_issue_detection", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.SUCCEEDED, result='{"summary": "Found 3 issues", "issues": 3, "total_traces_analyzed": 10}', retry_count=0, last_update_time=1234567900000, status_details=None, ) with ( mock.patch("mlflow.server.jobs.get_job", return_value=mock_job), app.test_client() as c, ): resp = c.get("/ajax-api/3.0/mlflow/jobs/job-123") assert resp.status_code == 200 json_response = resp.get_json() assert json_response["status"] == "SUCCEEDED" assert json_response["result"]["summary"] == "Found 3 issues" assert json_response["result"]["issues"] == 3 assert json_response["result"]["total_traces_analyzed"] == 10 assert json_response["status_details"] is None def test_get_job_pending(mock_job_store): mock_job = JobEntity( job_id="job-pending", creation_time=1234567890000, job_name="invoke_issue_detection", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.PENDING, result=None, retry_count=0, last_update_time=1234567890000, status_details=None, ) with ( mock.patch("mlflow.server.jobs.get_job", return_value=mock_job), app.test_client() as c, ): resp = c.get("/ajax-api/3.0/mlflow/jobs/job-pending") assert resp.status_code == 200 json_response = resp.get_json() assert json_response["status"] == "PENDING" assert json_response["result"] is None assert json_response["status_details"] is None def test_cancel_job_success(mock_job_store): mock_job = JobEntity( job_id="job-123", creation_time=1234567890000, job_name="invoke_issue_detection", params='{"experiment_id": "exp-123"}', timeout=None, status=JobStatus.CANCELED, result=None, retry_count=0, last_update_time=1234567900000, status_details=None, ) with ( mock.patch("mlflow.server.jobs.cancel_job", return_value=mock_job) as mock_cancel, app.test_client() as c, ): resp = c.patch("/ajax-api/3.0/mlflow/jobs/cancel/job-123") assert resp.status_code == 200 json_response = resp.get_json() assert json_response["status"] == "CANCELED" mock_cancel.assert_called_once_with("job-123") def test_get_rest_path_respects_static_prefix(monkeypatch): # Without prefix, both return bare paths assert _get_rest_path("/mlflow/experiments/search") == "/api/2.0/mlflow/experiments/search" assert _get_ajax_path("/mlflow/experiments/search") == "/ajax-api/2.0/mlflow/experiments/search" # With prefix, both should include the prefix monkeypatch.setenv(STATIC_PREFIX_ENV_VAR, "/myapp") assert ( _get_rest_path("/mlflow/experiments/search") == "/myapp/api/2.0/mlflow/experiments/search" ) assert ( _get_ajax_path("/mlflow/experiments/search") == "/myapp/ajax-api/2.0/mlflow/experiments/search" ) def _review_queue_entity(**overrides): defaults = { "queue_id": "rq-1", "experiment_id": "exp-1", "name": "q", "queue_type": ReviewQueueType.CUSTOM, "created_by": "alice", "creation_time_ms": 1, "last_update_time_ms": 1, "users": [], "schema_ids": [], } defaults.update(overrides) return ReviewQueue(**defaults) def test_create_review_queue_stamps_owner_from_authenticated_user(): request_message = CreateReviewQueue() request_message.experiment_id = "exp-1" request_message.name = "q" request_message.queue_type = ReviewQueueType.CUSTOM.to_proto() # The client even tries to set a different owner; it must be ignored. request_message.created_by = "spoofed" with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), mock.patch("mlflow.server.handlers._get_request_username", return_value="alice"), ): mock_store.return_value.create_review_queue.return_value = _review_queue_entity() _create_review_queue() call_kwargs = mock_store.return_value.create_review_queue.call_args[1] assert call_kwargs["created_by"] == "alice" def test_create_review_queue_no_owner_on_noauth(): request_message = CreateReviewQueue() request_message.experiment_id = "exp-1" request_message.name = "q" request_message.queue_type = ReviewQueueType.CUSTOM.to_proto() with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), # No auth plugin -> no request user. mock.patch("mlflow.server.handlers._get_request_username", return_value=None), ): mock_store.return_value.create_review_queue.return_value = _review_queue_entity( created_by=None ) _create_review_queue() call_kwargs = mock_store.return_value.create_review_queue.call_args[1] assert "created_by" not in call_kwargs def test_update_review_queue_passes_name_and_new_owner(): request_message = UpdateReviewQueue() request_message.queue_id = "rq-1" request_message.name = "renamed" request_message.new_owner = "bob" with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.update_review_queue.return_value = _review_queue_entity( name="renamed", created_by="bob" ) _update_review_queue() call_kwargs = mock_store.return_value.update_review_queue.call_args[1] assert call_kwargs["name"] == "renamed" assert call_kwargs["new_owner"] == "bob" # Untouched association sets stay None (proto2 presence on the singular # fields; no update_* flag set for users/schemas). assert call_kwargs["users"] is None assert call_kwargs["schema_ids"] is None def test_update_review_queue_unset_name_and_owner_pass_none(): # Only users changed; name / new_owner are unset, so HasField is False and # the store must receive None (guards against a truthiness-vs-HasField # regression that could, e.g., wipe the owner with an empty string). request_message = UpdateReviewQueue() request_message.queue_id = "rq-1" request_message.update_users = True request_message.users.append("bob") with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.update_review_queue.return_value = _review_queue_entity( users=["bob"] ) _update_review_queue() call_kwargs = mock_store.return_value.update_review_queue.call_args[1] assert call_kwargs["name"] is None assert call_kwargs["new_owner"] is None assert call_kwargs["users"] == ["bob"] def test_get_or_create_user_queue_ignores_client_created_by(): request_message = GetOrCreateUserQueue() request_message.experiment_id = "exp-1" request_message.user = "alice" request_message.created_by = "spoofed" with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), ): mock_store.return_value.get_or_create_user_queue.return_value = _review_queue_entity( queue_type=ReviewQueueType.USER, name="alice", created_by="alice", users=["alice"] ) _get_or_create_user_queue() call_kwargs = mock_store.return_value.get_or_create_user_queue.call_args[1] assert "created_by" not in call_kwargs assert call_kwargs["user"] == "alice" def _review_queue_item_entity(**overrides): defaults = { "queue_id": "rq-1", "item_type": ReviewItemType.TRACE, "item_id": "tr-1", "status": ReviewStatus.COMPLETE, "creation_time_ms": 1, "last_update_time_ms": 1, "completed_by": "alice", "completed_time_ms": 1, } defaults.update(overrides) return ReviewQueueItem(**defaults) def _set_status_request(status, completed_by=None): request_message = SetReviewQueueItemStatus() request_message.queue_id = "rq-1" request_message.item_id = "tr-1" request_message.status = status.to_proto() if completed_by is not None: request_message.completed_by = completed_by return request_message @pytest.mark.parametrize( ("username", "status", "client_completed_by", "expected_completed_by"), [ # Auth server: the caller's identity is stamped and the client value (even # a spoofed one) is ignored, for both terminal states. ("alice", ReviewStatus.COMPLETE, "ghost@nowhere.example", "alice"), ("alice", ReviewStatus.DECLINED, "ghost@nowhere.example", "alice"), # Auth server reopen: attribution is cleared regardless of the client value. ("alice", ReviewStatus.PENDING, "ghost@nowhere.example", None), # No-auth server: no identity to bind to, so the client value passes through # verbatim (a real no-auth client sends the single `default` user). (None, ReviewStatus.COMPLETE, "default", "default"), (None, ReviewStatus.COMPLETE, "ghost@nowhere.example", "ghost@nowhere.example"), ], ) def test_set_review_queue_item_status_stamps_completed_by( username, status, client_completed_by, expected_completed_by ): request_message = _set_status_request(status, completed_by=client_completed_by) with ( mock.patch("mlflow.server.handlers._get_tracking_store") as mock_store, mock.patch("mlflow.server.handlers._get_request_message", return_value=request_message), mock.patch("mlflow.server.handlers._get_request_username", return_value=username), ): mock_store.return_value.set_review_queue_item_status.return_value = ( _review_queue_item_entity( status=status, completed_by=expected_completed_by, completed_time_ms=None if status == ReviewStatus.PENDING else 1, ) ) _set_review_queue_item_status() mock_store.return_value.set_review_queue_item_status.assert_called_once() call_kwargs = mock_store.return_value.set_review_queue_item_status.call_args[1] assert call_kwargs["completed_by"] == expected_completed_by # Pin status pass-through too, so a regression that mangles status is caught. assert call_kwargs["status"] == status