1  from pathlib import Path
  2  
  3  import pytest
  4  import requests
  5  
  6  import mlflow
  7  from mlflow.server import handlers
  8  from mlflow.server.fastapi_app import app
  9  from mlflow.server.handlers import initialize_backend_stores
 10  
 11  from tests.helper_functions import get_safe_port
 12  from tests.tracking.integration_test_utils import ServerThread
 13  
 14  
 15  @pytest.fixture
 16  def tracking_server(db_uri: str, tmp_path: Path):
 17      handlers._tracking_store = None
 18      handlers._model_registry_store = None
 19      initialize_backend_stores(db_uri, default_artifact_root=tmp_path.as_uri())
 20  
 21      with ServerThread(app, get_safe_port()) as url:
 22          mlflow.set_tracking_uri(url)
 23          yield url
 24          mlflow.set_tracking_uri(None)
 25  
 26  
 27  def test_generate_demo_route_creates_data(tracking_server):
 28      response = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 29  
 30      assert response.status_code == 200
 31      data = response.json()
 32      assert data["status"] == "created"
 33      assert data["experiment_id"] is not None
 34      assert len(data["features_generated"]) > 0
 35      assert data["navigation_url"] is not None
 36  
 37  
 38  def test_generate_demo_route_idempotent(tracking_server):
 39      response1 = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 40      assert response1.status_code == 200
 41      data1 = response1.json()
 42      assert data1["status"] == "created"
 43  
 44      response2 = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 45      assert response2.status_code == 200
 46      data2 = response2.json()
 47      assert data2["status"] == "exists"
 48      assert data2["features_generated"] == []
 49  
 50  
 51  def test_delete_demo_route_removes_data(tracking_server):
 52      requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 53  
 54      response = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/delete")
 55  
 56      assert response.status_code == 200
 57      data = response.json()
 58      assert data["status"] == "deleted"
 59      assert len(data["features_deleted"]) > 0
 60  
 61  
 62  def test_delete_demo_route_when_no_data(tracking_server):
 63      response = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/delete")
 64  
 65      assert response.status_code == 200
 66      data = response.json()
 67      assert data["status"] == "deleted"
 68      assert data["features_deleted"] == []
 69  
 70  
 71  def test_generate_after_delete_regenerates(tracking_server):
 72      response1 = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 73      assert response1.json()["status"] == "created"
 74  
 75      requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/delete")
 76  
 77      response2 = requests.post(f"{tracking_server}/ajax-api/3.0/mlflow/demo/generate")
 78      assert response2.status_code == 200
 79      data = response2.json()
 80      assert data["status"] == "created"
 81      assert len(data["features_generated"]) > 0
 82  
 83  
 84  def test_get_demo_endpoints_returns_routes():
 85      from mlflow.server.handlers import get_demo_endpoints
 86  
 87      endpoints = get_demo_endpoints()
 88  
 89      assert len(endpoints) == 2
 90  
 91      paths = [path for path, _, _ in endpoints]
 92      assert any("/mlflow/demo/generate" in path for path in paths)
 93      assert any("/mlflow/demo/delete" in path for path in paths)
 94  
 95  
 96  def test_get_endpoints_includes_demo():
 97      from mlflow.server.handlers import get_endpoints
 98  
 99      endpoints = get_endpoints()
100      paths = [path for path, _, _ in endpoints]
101  
102      assert any("/mlflow/demo/generate" in path for path in paths)
103      assert any("/mlflow/demo/delete" in path for path in paths)