wehub-resource-sync
493 lines
18 KiB
Python
493 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
from mlflow.entities.workspace import Workspace, WorkspaceDeletionMode
|
|
from mlflow.server.handlers import get_endpoints
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def enable_workspaces(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_ENABLE_WORKSPACES", "true")
|
|
|
|
|
|
@pytest.fixture
|
|
def app(monkeypatch):
|
|
flask_app = Flask(__name__)
|
|
for rule, view_func, methods in get_endpoints():
|
|
flask_app.add_url_rule(rule, view_func=view_func, methods=methods)
|
|
return flask_app
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_workspace_store(monkeypatch):
|
|
store = mock.Mock()
|
|
monkeypatch.setattr(
|
|
"mlflow.server.handlers._get_workspace_store",
|
|
lambda *_, **__: store,
|
|
)
|
|
return store
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_tracking_store(monkeypatch):
|
|
store = mock.Mock()
|
|
store.artifact_root_uri = "/default/artifact/root"
|
|
monkeypatch.setattr(
|
|
"mlflow.server.handlers._get_tracking_store",
|
|
lambda *_, **__: store,
|
|
)
|
|
return store
|
|
|
|
|
|
def _workspace_to_json(payload):
|
|
return json.loads(payload)
|
|
|
|
|
|
def test_list_workspaces_endpoint(app, mock_workspace_store):
|
|
mock_workspace_store.list_workspaces.return_value = [
|
|
Workspace(name="default", description="Default"),
|
|
Workspace(name="team-a", description=None),
|
|
]
|
|
with app.test_client() as client:
|
|
response = client.get("/api/3.0/mlflow/workspaces")
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["workspaces"][0] == {"name": "default", "description": "Default"}
|
|
assert payload["workspaces"][1] == {"name": "team-a"}
|
|
mock_workspace_store.list_workspaces.assert_called_once_with()
|
|
|
|
|
|
def test_create_workspace_endpoint(app, mock_workspace_store, mock_tracking_store):
|
|
created = Workspace(
|
|
name="team-b",
|
|
description="Team B",
|
|
trace_archival_location="s3://archive/team-b",
|
|
trace_archival_retention="30d",
|
|
)
|
|
mock_workspace_store.create_workspace.return_value = created
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={
|
|
"name": "team-b",
|
|
"description": "Team B",
|
|
"trace_archival_config": {
|
|
"location": "s3://archive/team-b",
|
|
"retention": "30d",
|
|
},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 201
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {
|
|
"workspace": {
|
|
"name": "team-b",
|
|
"description": "Team B",
|
|
"trace_archival_config": {
|
|
"location": "s3://archive/team-b",
|
|
"retention": "30d",
|
|
},
|
|
}
|
|
}
|
|
mock_workspace_store.create_workspace.assert_called_once()
|
|
args, _ = mock_workspace_store.create_workspace.call_args
|
|
assert args[0].trace_archival_location == "s3://archive/team-b"
|
|
assert args[0].trace_archival_retention == "30d"
|
|
mock_tracking_store.get_experiment_by_name.assert_not_called()
|
|
mock_tracking_store.create_experiment.assert_not_called()
|
|
|
|
|
|
def test_get_workspace_endpoint(app, mock_workspace_store):
|
|
mock_workspace_store.get_workspace.return_value = Workspace(name="team-c", description="Team C")
|
|
with app.test_client() as client:
|
|
response = client.get("/api/3.0/mlflow/workspaces/team-c")
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {"workspace": {"name": "team-c", "description": "Team C"}}
|
|
mock_workspace_store.get_workspace.assert_called_once_with("team-c")
|
|
|
|
|
|
def test_update_workspace_endpoint(app, mock_workspace_store):
|
|
updated = Workspace(
|
|
name="team-d",
|
|
description="Updated",
|
|
trace_archival_location="s3://archive/team-d",
|
|
trace_archival_retention="14d",
|
|
)
|
|
mock_workspace_store.update_workspace.return_value = updated
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-d",
|
|
json={
|
|
"description": "Updated",
|
|
"trace_archival_config": {
|
|
"location": "s3://archive/team-d",
|
|
"retention": "14d",
|
|
},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {
|
|
"workspace": {
|
|
"name": "team-d",
|
|
"description": "Updated",
|
|
"trace_archival_config": {
|
|
"location": "s3://archive/team-d",
|
|
"retention": "14d",
|
|
},
|
|
}
|
|
}
|
|
mock_workspace_store.update_workspace.assert_called_once()
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert args[0].trace_archival_location == "s3://archive/team-d"
|
|
assert args[0].trace_archival_retention == "14d"
|
|
|
|
|
|
def test_update_default_workspace_allows_reserved_name(app, mock_workspace_store):
|
|
updated = Workspace(name="default", default_artifact_root="s3://bucket/root")
|
|
mock_workspace_store.update_workspace.return_value = updated
|
|
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/default",
|
|
json={"default_artifact_root": "s3://bucket/root"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {
|
|
"workspace": {"name": "default", "default_artifact_root": "s3://bucket/root"}
|
|
}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert args[0].name == "default"
|
|
assert args[0].default_artifact_root == "s3://bucket/root"
|
|
|
|
|
|
def test_update_workspace_can_clear_default_artifact_root(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
cleared = Workspace(name="team-clear", description=None, default_artifact_root=None)
|
|
mock_workspace_store.update_workspace.return_value = cleared
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-clear",
|
|
json={"default_artifact_root": " "},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {"workspace": {"name": "team-clear"}}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert isinstance(args[0], Workspace)
|
|
assert args[0].name == "team-clear"
|
|
# Handler passes "" to indicate "clear"; the store converts "" to None
|
|
assert args[0].default_artifact_root == ""
|
|
|
|
|
|
def test_update_workspace_can_clear_trace_archival_location(app, mock_workspace_store):
|
|
cleared = Workspace(name="team-clear", description=None, trace_archival_location=None)
|
|
mock_workspace_store.update_workspace.return_value = cleared
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-clear",
|
|
json={"trace_archival_config": {"location": ""}},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {"workspace": {"name": "team-clear"}}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert isinstance(args[0], Workspace)
|
|
assert args[0].name == "team-clear"
|
|
assert args[0].trace_archival_location == ""
|
|
|
|
|
|
def test_update_workspace_can_clear_trace_archival_retention(app, mock_workspace_store):
|
|
cleared = Workspace(name="team-clear", description=None, trace_archival_retention=None)
|
|
mock_workspace_store.update_workspace.return_value = cleared
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-clear",
|
|
json={"trace_archival_config": {"retention": ""}},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {"workspace": {"name": "team-clear"}}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert isinstance(args[0], Workspace)
|
|
assert args[0].name == "team-clear"
|
|
assert args[0].trace_archival_retention == ""
|
|
|
|
|
|
def test_update_workspace_with_only_trace_archival_location(app, mock_workspace_store):
|
|
updated = Workspace(name="team-only-location", trace_archival_location="s3://archive/team-only")
|
|
mock_workspace_store.update_workspace.return_value = updated
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-only-location",
|
|
json={"trace_archival_config": {"location": "s3://archive/team-only"}},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {
|
|
"workspace": {
|
|
"name": "team-only-location",
|
|
"trace_archival_config": {"location": "s3://archive/team-only"},
|
|
}
|
|
}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert args[0].trace_archival_location == "s3://archive/team-only"
|
|
assert args[0].trace_archival_retention is None
|
|
|
|
|
|
def test_update_workspace_with_only_trace_archival_retention(app, mock_workspace_store):
|
|
updated = Workspace(name="team-only-retention", trace_archival_retention="14d")
|
|
mock_workspace_store.update_workspace.return_value = updated
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-only-retention",
|
|
json={"trace_archival_config": {"retention": "14d"}},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload == {
|
|
"workspace": {
|
|
"name": "team-only-retention",
|
|
"trace_archival_config": {"retention": "14d"},
|
|
}
|
|
}
|
|
args, _ = mock_workspace_store.update_workspace.call_args
|
|
assert args[0].trace_archival_location is None
|
|
assert args[0].trace_archival_retention == "14d"
|
|
|
|
|
|
def test_create_workspace_rejects_invalid_trace_archival_retention(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={"name": "team-bad", "trace_archival_config": {"retention": "90days"}},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.retention'. Expected a duration in the form "
|
|
"`<int><unit>`, where unit is one of 'm', 'h', or 'd' (for example '30d' or '12h')."
|
|
)
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_create_workspace_rejects_invalid_trace_archival_location(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={
|
|
"name": "team-bad",
|
|
"trace_archival_config": {"location": "s3://archive/team#fragment"},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert "trace_archival_config.location" in payload["message"]
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_create_workspace_rejects_proxy_only_trace_archival_location(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={
|
|
"name": "team-proxy",
|
|
"trace_archival_config": {"location": "mlflow-artifacts:/archive/team-proxy"},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.location'. Trace archival location cannot use "
|
|
"the proxy-only `mlflow-artifacts:` scheme."
|
|
)
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_create_workspace_rejects_non_uri_trace_archival_location(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={
|
|
"name": "team-local-path",
|
|
"trace_archival_config": {"location": "archive/team-local-path"},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.location'. Expected a URI string."
|
|
)
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_update_workspace_rejects_invalid_trace_archival_retention(app, mock_workspace_store):
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-bad",
|
|
json={"trace_archival_config": {"retention": "5w"}},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.retention'. Expected a duration in the form "
|
|
"`<int><unit>`, where unit is one of 'm', 'h', or 'd' (for example '30d' or '12h')."
|
|
)
|
|
mock_workspace_store.update_workspace.assert_not_called()
|
|
|
|
|
|
def test_create_workspace_rejects_trace_archival_retention_longer_than_32_chars(
|
|
app, mock_workspace_store, mock_tracking_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={"name": "team-bad", "trace_archival_config": {"retention": f"{'1' * 32}d"}},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.retention'. Maximum length is 32 characters."
|
|
)
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_update_workspace_rejects_trace_archival_retention_longer_than_32_chars(
|
|
app, mock_workspace_store
|
|
):
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-bad",
|
|
json={"trace_archival_config": {"retention": f"{'1' * 32}d"}},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert payload["message"] == (
|
|
"Invalid value for 'trace_archival_config.retention'. Maximum length is 32 characters."
|
|
)
|
|
mock_workspace_store.update_workspace.assert_not_called()
|
|
|
|
|
|
def test_delete_workspace_endpoint(app, mock_workspace_store):
|
|
with app.test_client() as client:
|
|
response = client.delete("/api/3.0/mlflow/workspaces/team-e")
|
|
|
|
assert response.status_code == 204
|
|
mock_workspace_store.delete_workspace.assert_called_once_with(
|
|
"team-e", mode=WorkspaceDeletionMode.RESTRICT
|
|
)
|
|
|
|
|
|
def test_delete_default_workspace_rejected_by_validation(app, mock_workspace_store):
|
|
with app.test_client() as client:
|
|
response = client.delete("/api/3.0/mlflow/workspaces/default")
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert "cannot be deleted" in payload["message"]
|
|
mock_workspace_store.delete_workspace.assert_not_called()
|
|
|
|
|
|
def test_create_workspace_fails_without_artifact_root(app, mock_workspace_store, monkeypatch):
|
|
tracking_store = mock.Mock()
|
|
tracking_store.artifact_root_uri = None
|
|
monkeypatch.setattr(
|
|
"mlflow.server.handlers._get_tracking_store",
|
|
lambda *_, **__: tracking_store,
|
|
)
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={"name": "team-no-root"},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert "artifact root" in payload["message"].lower()
|
|
|
|
|
|
def test_create_workspace_with_artifact_root_succeeds_without_server_default(
|
|
app, mock_workspace_store, monkeypatch
|
|
):
|
|
tracking_store = mock.Mock()
|
|
tracking_store.artifact_root_uri = None
|
|
monkeypatch.setattr(
|
|
"mlflow.server.handlers._get_tracking_store",
|
|
lambda *_, **__: tracking_store,
|
|
)
|
|
created = Workspace(name="team-with-root", default_artifact_root="s3://bucket/path")
|
|
mock_workspace_store.create_workspace.return_value = created
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={"name": "team-with-root", "default_artifact_root": "s3://bucket/path"},
|
|
)
|
|
|
|
assert response.status_code == 201
|
|
|
|
|
|
def test_create_default_workspace_rejected(app, mock_workspace_store, mock_tracking_store):
|
|
with app.test_client() as client:
|
|
response = client.post(
|
|
"/api/3.0/mlflow/workspaces",
|
|
json={"name": "default"},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert "reserved" in payload["message"]
|
|
mock_workspace_store.create_workspace.assert_not_called()
|
|
|
|
|
|
def test_update_workspace_clear_artifact_root_fails_without_server_default(
|
|
app, mock_workspace_store, monkeypatch
|
|
):
|
|
tracking_store = mock.Mock()
|
|
tracking_store.artifact_root_uri = None
|
|
monkeypatch.setattr(
|
|
"mlflow.server.handlers._get_tracking_store",
|
|
lambda *_, **__: tracking_store,
|
|
)
|
|
with app.test_client() as client:
|
|
response = client.patch(
|
|
"/api/3.0/mlflow/workspaces/team-clear",
|
|
json={"default_artifact_root": ""},
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
payload = _workspace_to_json(response.get_data(True))
|
|
assert "artifact root" in payload["message"].lower()
|