python_integration_test.py

#!/usr/bin/env python3
"""
Integration test script to verify MLflow Model Registry Rust API compatibility with Python MLflow.

This script tests the Rust implementation against a running Rust server
to ensure API compatibility with the Python MLflow client.
"""

import json
import requests
import time
import sys
from typing import Optional, Dict, Any

class MLflowRegistryRustClient:
    """Client for testing the Rust MLflow Registry API"""

    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = requests.Session()

    def _make_request(self, method: str, endpoint: str, data: Optional[Dict[Any, Any]] = None) -> Dict[Any, Any]:
        url = f"{self.base_url}{endpoint}"

        if method.upper() == "GET":
            response = self.session.get(url, params=data)
        elif method.upper() == "POST":
            response = self.session.post(url, json=data)
        elif method.upper() == "PATCH":
            response = self.session.patch(url, json=data)
        elif method.upper() == "DELETE":
            response = self.session.delete(url, json=data)
        else:
            raise ValueError(f"Unsupported HTTP method: {method}")

        if response.status_code >= 400:
            print(f"Error {response.status_code}: {response.text}")
            response.raise_for_status()

        if response.status_code == 204:  # No content
            return {}

        return response.json()

    # Registered Model operations
    def create_registered_model(self, name: str, description: Optional[str] = None,
                              tags: Optional[Dict[str, str]] = None) -> Dict[Any, Any]:
        data = {"name": name}
        if description:
            data["description"] = description
        if tags:
            data["tags"] = tags

        return self._make_request("POST", "/api/2.0/mlflow/model-registry/registered-models/create", data)

    def get_registered_model(self, name: str) -> Dict[Any, Any]:
        return self._make_request("GET", f"/models/{name}")

    def search_registered_models(self, filter_string: Optional[str] = None,
                               max_results: Optional[int] = None) -> Dict[Any, Any]:
        params = {}
        if filter_string:
            params["filter_string"] = filter_string
        if max_results:
            params["max_results"] = max_results

        return self._make_request("GET", "/api/2.0/mlflow/model-registry/registered-models/search", params)

    def update_registered_model(self, name: str, description: Optional[str] = None) -> Dict[Any, Any]:
        data = {"name": name}
        if description:
            data["description"] = description

        return self._make_request("PATCH", "/api/2.0/mlflow/model-registry/registered-models/update", data)

    def delete_registered_model(self, name: str) -> None:
        self._make_request("DELETE", f"/models/{name}")

    # Model Version operations
    def create_model_version(self, name: str, source: str, run_id: Optional[str] = None,
                           description: Optional[str] = None, tags: Optional[Dict[str, str]] = None) -> Dict[Any, Any]:
        data = {"name": name, "source": source}
        if run_id:
            data["run_id"] = run_id
        if description:
            data["description"] = description
        if tags:
            data["tags"] = tags

        return self._make_request("POST", "/api/2.0/mlflow/model-registry/model-versions/create", data)

    def get_model_version(self, name: str, version: str) -> Dict[Any, Any]:
        return self._make_request("GET", f"/models/{name}/versions/{version}")

    def search_model_versions(self, filter_string: Optional[str] = None,
                            max_results: Optional[int] = None) -> Dict[Any, Any]:
        params = {}
        if filter_string:
            params["filter_string"] = filter_string
        if max_results:
            params["max_results"] = max_results

        return self._make_request("GET", "/api/2.0/mlflow/model-registry/model-versions/search", params)

    def update_model_version(self, name: str, version: str, description: Optional[str] = None) -> Dict[Any, Any]:
        data = {"name": name, "version": version}
        if description:
            data["description"] = description

        return self._make_request("PATCH", "/api/2.0/mlflow/model-registry/model-versions/update", data)

    def delete_model_version(self, name: str, version: str) -> None:
        self._make_request("DELETE", f"/models/{name}/versions/{version}")

    def transition_model_version_stage(self, name: str, version: str, stage: str,
                                     archive_existing_versions: bool = False) -> Dict[Any, Any]:
        data = {
            "name": name,
            "version": version,
            "stage": stage,
            "archive_existing_versions": archive_existing_versions
        }
        return self._make_request("POST", "/api/2.0/mlflow/model-registry/model-versions/transition-stage", data)

    def get_model_version_download_uri(self, name: str, version: str) -> str:
        response = self._make_request("GET", "/api/2.0/mlflow/model-registry/model-versions/get-download-uri",
                                    {"name": name, "version": version})
        return response["artifact_uri"]

    # Alias operations
    def set_registered_model_alias(self, name: str, alias: str, version: str) -> None:
        data = {"name": name, "alias": alias, "version": version}
        self._make_request("POST", "/api/2.0/mlflow/model-registry/registered-models/set-alias", data)

    def delete_registered_model_alias(self, name: str, alias: str) -> None:
        data = {"name": name, "alias": alias}
        self._make_request("DELETE", "/api/2.0/mlflow/model-registry/registered-models/delete-alias", data)

    def get_model_version_by_alias(self, name: str, alias: str) -> Dict[Any, Any]:
        return self._make_request("GET", f"/models/{name}/aliases/{alias}")

    # Tag operations
    def set_registered_model_tag(self, name: str, key: str, value: str) -> None:
        data = {"name": name, "key": key, "value": value}
        self._make_request("POST", "/api/2.0/mlflow/model-registry/registered-models/set-tag", data)

    def delete_registered_model_tag(self, name: str, key: str) -> None:
        data = {"name": name, "key": key}
        self._make_request("DELETE", "/api/2.0/mlflow/model-registry/registered-models/delete-tag", data)

    def set_model_version_tag(self, name: str, version: str, key: str, value: str) -> None:
        data = {"name": name, "version": version, "key": key, "value": value}
        self._make_request("POST", "/api/2.0/mlflow/model-registry/model-versions/set-tag", data)

    def delete_model_version_tag(self, name: str, version: str, key: str) -> None:
        data = {"name": name, "version": version, "key": key}
        self._make_request("DELETE", "/api/2.0/mlflow/model-registry/model-versions/delete-tag", data)


def test_basic_model_lifecycle(client: MLflowRegistryRustClient):
    """Test basic model lifecycle: create, version, stage transitions"""
    print("Testing basic model lifecycle...")

    model_name = f"python_test_model_{int(time.time())}"

    # Create registered model
    created_model = client.create_registered_model(
        name=model_name,
        description="Test model created from Python",
        tags={"test": "true", "language": "python"}
    )

    assert created_model["name"] == model_name
    assert created_model["description"] == "Test model created from Python"
    print(f"✓ Created model: {model_name}")

    # Get registered model
    retrieved_model = client.get_registered_model(model_name)
    assert retrieved_model["name"] == model_name
    assert retrieved_model["tags"]["test"] == "true"
    assert retrieved_model["tags"]["language"] == "python"
    print(f"✓ Retrieved model: {model_name}")

    # Create model version
    version1 = client.create_model_version(
        name=model_name,
        source="s3://test-bucket/python-test/v1/",
        run_id="python-run-123",
        description="First version from Python",
        tags={"accuracy": "0.95", "framework": "scikit-learn"}
    )

    assert version1["name"] == model_name
    assert version1["version"] == "1"
    assert version1["run_id"] == "python-run-123"
    print(f"✓ Created version 1 for model: {model_name}")

    # Create second version
    version2 = client.create_model_version(
        name=model_name,
        source="s3://test-bucket/python-test/v2/",
        run_id="python-run-124",
        description="Second version from Python",
        tags={"accuracy": "0.97", "framework": "pytorch"}
    )

    assert version2["version"] == "2"
    print(f"✓ Created version 2 for model: {model_name}")

    # Test stage transitions
    staging_version = client.transition_model_version_stage(model_name, "1", "staging")
    assert staging_version["current_stage"] == "staging"
    print(f"✓ Transitioned version 1 to staging")

    production_version = client.transition_model_version_stage(model_name, "2", "production")
    assert production_version["current_stage"] == "production"
    print(f"✓ Transitioned version 2 to production")

    # Test aliases
    client.set_registered_model_alias(model_name, "champion", "2")
    client.set_registered_model_alias(model_name, "challenger", "1")

    champion = client.get_model_version_by_alias(model_name, "champion")
    assert champion["version"] == "2"

    challenger = client.get_model_version_by_alias(model_name, "challenger")
    assert challenger["version"] == "1"
    print(f"✓ Set and verified aliases")

    # Test download URI
    download_uri = client.get_model_version_download_uri(model_name, "1")
    assert "s3://" in download_uri
    print(f"✓ Got download URI: {download_uri}")

    # Cleanup
    client.delete_registered_model(model_name)
    print(f"✓ Cleaned up model: {model_name}")


def test_tag_operations(client: MLflowRegistryRustClient):
    """Test tag operations on models and versions"""
    print("Testing tag operations...")

    model_name = f"python_tag_test_{int(time.time())}"

    # Create model and version
    client.create_registered_model(model_name)
    version = client.create_model_version(model_name, "s3://test/tag-test/")

    # Test model tags
    client.set_registered_model_tag(model_name, "env", "production")
    client.set_registered_model_tag(model_name, "owner", "ml-team")

    model_with_tags = client.get_registered_model(model_name)
    assert model_with_tags["tags"]["env"] == "production"
    assert model_with_tags["tags"]["owner"] == "ml-team"
    print("✓ Set and verified model tags")

    # Test version tags
    client.set_model_version_tag(model_name, version["version"], "accuracy", "0.98")
    client.set_model_version_tag(model_name, version["version"], "validated", "true")

    version_with_tags = client.get_model_version(model_name, version["version"])
    assert version_with_tags["tags"]["accuracy"] == "0.98"
    assert version_with_tags["tags"]["validated"] == "true"
    print("✓ Set and verified version tags")

    # Test tag deletion
    client.delete_registered_model_tag(model_name, "owner")
    client.delete_model_version_tag(model_name, version["version"], "validated")

    model_after_delete = client.get_registered_model(model_name)
    version_after_delete = client.get_model_version(model_name, version["version"])

    assert "owner" not in model_after_delete["tags"]
    assert "validated" not in version_after_delete["tags"]
    assert model_after_delete["tags"]["env"] == "production"
    assert version_after_delete["tags"]["accuracy"] == "0.98"
    print("✓ Verified tag deletion")

    # Cleanup
    client.delete_registered_model(model_name)


def test_search_operations(client: MLflowRegistryRustClient):
    """Test search operations"""
    print("Testing search operations...")

    # Create multiple models
    models = []
    for i in range(3):
        model_name = f"python_search_test_{i}_{int(time.time())}"
        client.create_registered_model(
            model_name,
            f"Search test model {i}",
            {"search_test": "true", "index": str(i)}
        )
        models.append(model_name)

        # Create versions for each model
        client.create_model_version(model_name, f"s3://test/search/{i}/v1/")
        client.create_model_version(model_name, f"s3://test/search/{i}/v2/")

    # Search registered models
    search_results = client.search_registered_models(max_results=10)
    found_models = [m for m in search_results["registered_models"] if m["name"] in models]
    assert len(found_models) == 3
    print("✓ Found all created models in search")

    # Search model versions
    version_results = client.search_model_versions(max_results=20)
    found_versions = [v for v in version_results["model_versions"] if v["name"] in models]
    assert len(found_versions) >= 6  # 2 versions per model * 3 models
    print("✓ Found all created versions in search")

    # Test with max_results limit
    limited_results = client.search_registered_models(max_results=2)
    assert len(limited_results["registered_models"]) <= 2
    print("✓ Verified max_results limiting")

    # Cleanup
    for model_name in models:
        client.delete_registered_model(model_name)


def test_error_handling(client: MLflowRegistryRustClient):
    """Test error handling for invalid operations"""
    print("Testing error handling...")

    # Try to get non-existent model
    try:
        client.get_registered_model("non_existent_model_12345")
        assert False, "Should have raised an error"
    except requests.exceptions.HTTPError as e:
        assert e.response.status_code == 404
        print("✓ Correctly handled non-existent model")

    # Try to create duplicate model
    model_name = f"python_duplicate_test_{int(time.time())}"
    client.create_registered_model(model_name)

    try:
        client.create_registered_model(model_name)
        assert False, "Should have raised an error for duplicate model"
    except requests.exceptions.HTTPError as e:
        assert e.response.status_code == 409
        print("✓ Correctly handled duplicate model creation")

    # Try to get non-existent version
    try:
        client.get_model_version(model_name, "999")
        assert False, "Should have raised an error"
    except requests.exceptions.HTTPError as e:
        assert e.response.status_code == 404
        print("✓ Correctly handled non-existent version")

    # Cleanup
    client.delete_registered_model(model_name)


def wait_for_server(base_url: str, timeout: int = 30):
    """Wait for the server to be ready"""
    print(f"Waiting for server at {base_url}...")

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            response = requests.get(f"{base_url}/health", timeout=5)
            if response.status_code == 200:
                print("✓ Server is ready")
                return True
        except requests.exceptions.RequestException:
            pass

        time.sleep(1)

    print(f"✗ Server not ready after {timeout} seconds")
    return False


def main():
    """Run all integration tests"""
    base_url = "http://localhost:8000"

    if not wait_for_server(base_url):
        print("Server is not running. Please start the Rust MLflow Registry server first.")
        sys.exit(1)

    client = MLflowRegistryRustClient(base_url)

    tests = [
        test_basic_model_lifecycle,
        test_tag_operations,
        test_search_operations,
        test_error_handling,
    ]

    print(f"Running {len(tests)} integration tests...\n")

    for i, test_func in enumerate(tests, 1):
        try:
            print(f"[{i}/{len(tests)}] {test_func.__name__}")
            test_func(client)
            print(f"✓ {test_func.__name__} passed\n")
        except Exception as e:
            print(f"✗ {test_func.__name__} failed: {e}\n")
            sys.exit(1)

    print("🎉 All integration tests passed!")


if __name__ == "__main__":
    main()