wehub-resource-sync
218 lines
8.7 KiB
Python
218 lines
8.7 KiB
Python
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import pytest
|
|
from cryptography.fernet import Fernet
|
|
|
|
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent, WebhookStatus
|
|
from mlflow.environment_variables import MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.server import handlers
|
|
from mlflow.server.fastapi_app import app
|
|
from mlflow.server.handlers import initialize_backend_stores
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
from tests.helper_functions import get_safe_port
|
|
from tests.tracking.integration_test_utils import ServerThread
|
|
|
|
|
|
@pytest.fixture
|
|
def client(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[MlflowClient]:
|
|
"""Setup a local MLflow server with proper webhook encryption key support."""
|
|
# Set up encryption key for webhooks using monkeypatch
|
|
encryption_key = Fernet.generate_key().decode("utf-8")
|
|
monkeypatch.setenv(MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY.name, encryption_key)
|
|
|
|
# Configure backend stores
|
|
backend_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
|
|
default_artifact_root = tmp_path.as_uri()
|
|
|
|
# Force-reset backend stores before each test
|
|
handlers._tracking_store = None
|
|
handlers._model_registry_store = None
|
|
initialize_backend_stores(backend_uri, default_artifact_root=default_artifact_root)
|
|
|
|
# Start server and return client
|
|
with ServerThread(app, get_safe_port()) as url:
|
|
yield MlflowClient(url)
|
|
|
|
|
|
def test_create_webhook(client: MlflowClient):
|
|
webhook = client.create_webhook(
|
|
name="test_webhook",
|
|
url="https://example.com/webhook",
|
|
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
|
|
)
|
|
assert webhook.name == "test_webhook"
|
|
assert webhook.url == "https://example.com/webhook"
|
|
assert webhook.secret is None
|
|
assert webhook.events == [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
|
|
|
|
webhook = client.get_webhook(webhook.webhook_id)
|
|
assert webhook.name == "test_webhook"
|
|
assert webhook.url == "https://example.com/webhook"
|
|
assert webhook.secret is None
|
|
|
|
# With secret
|
|
webhook_with_secret = client.create_webhook(
|
|
name="test_webhook_with_secret",
|
|
url="https://example.com/webhook_with_secret",
|
|
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
|
|
secret="my_secret",
|
|
)
|
|
assert webhook_with_secret.name == "test_webhook_with_secret"
|
|
assert webhook_with_secret.url == "https://example.com/webhook_with_secret"
|
|
assert webhook_with_secret.secret is None
|
|
assert webhook_with_secret.events == [
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)
|
|
]
|
|
|
|
# Multiple events
|
|
webhook_multiple_events = client.create_webhook(
|
|
name="test_webhook_multiple_events",
|
|
url="https://example.com/webhook_multiple_events",
|
|
events=[
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED),
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
|
|
],
|
|
)
|
|
assert webhook_multiple_events.name == "test_webhook_multiple_events"
|
|
assert webhook_multiple_events.url == "https://example.com/webhook_multiple_events"
|
|
assert sorted(
|
|
webhook_multiple_events.events, key=lambda e: (e.entity.value, e.action.value)
|
|
) == [
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED),
|
|
]
|
|
assert webhook_multiple_events.secret is None
|
|
|
|
|
|
def test_get_webhook(client: MlflowClient):
|
|
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
|
|
created_webhook = client.create_webhook(
|
|
name="test_webhook", url="https://example.com/webhook", events=events
|
|
)
|
|
retrieved_webhook = client.get_webhook(created_webhook.webhook_id)
|
|
assert retrieved_webhook.webhook_id == created_webhook.webhook_id
|
|
assert retrieved_webhook.name == "test_webhook"
|
|
assert retrieved_webhook.url == "https://example.com/webhook"
|
|
assert retrieved_webhook.events == events
|
|
|
|
|
|
def test_get_webhook_not_found(client: MlflowClient):
|
|
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
|
|
client.get_webhook("nonexistent")
|
|
|
|
|
|
def test_list_webhooks(client: MlflowClient):
|
|
# Create more webhooks than max_results
|
|
for i in range(5):
|
|
client.create_webhook(
|
|
name=f"webhook{i}",
|
|
url=f"https://example.com/{i}",
|
|
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
|
|
)
|
|
# Test pagination with max_results=2
|
|
webhooks_page = client.list_webhooks(max_results=2)
|
|
assert len(webhooks_page) == 2
|
|
assert webhooks_page.token is not None
|
|
# Get next page
|
|
next_webhooks_page = client.list_webhooks(max_results=2, page_token=webhooks_page.token)
|
|
assert len(next_webhooks_page) == 2
|
|
assert next_webhooks_page.token is not None
|
|
# Verify we don't get duplicates
|
|
first_page_ids = {w.webhook_id for w in webhooks_page}
|
|
second_page_ids = {w.webhook_id for w in next_webhooks_page}
|
|
assert first_page_ids.isdisjoint(second_page_ids)
|
|
|
|
|
|
def test_update_webhook(client: MlflowClient):
|
|
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
|
|
webhook = client.create_webhook(
|
|
name="original_name", url="https://example.com/original", events=events
|
|
)
|
|
# Update webhook
|
|
new_events = [
|
|
WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
|
|
WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED),
|
|
]
|
|
updated_webhook = client.update_webhook(
|
|
webhook_id=webhook.webhook_id,
|
|
name="updated_name",
|
|
url="https://example.com/updated",
|
|
events=new_events,
|
|
description="Updated description",
|
|
secret="new_secret",
|
|
status=WebhookStatus.DISABLED,
|
|
)
|
|
assert updated_webhook.webhook_id == webhook.webhook_id
|
|
assert updated_webhook.name == "updated_name"
|
|
assert updated_webhook.url == "https://example.com/updated"
|
|
assert updated_webhook.events == new_events
|
|
assert updated_webhook.description == "Updated description"
|
|
assert updated_webhook.status == WebhookStatus.DISABLED
|
|
assert updated_webhook.last_updated_timestamp > webhook.last_updated_timestamp
|
|
|
|
|
|
def test_update_webhook_partial(client: MlflowClient):
|
|
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
|
|
webhook = client.create_webhook(
|
|
name="original_name",
|
|
url="https://example.com/original",
|
|
events=events,
|
|
description="Original description",
|
|
)
|
|
# Update only the name
|
|
updated_webhook = client.update_webhook(
|
|
webhook_id=webhook.webhook_id,
|
|
name="updated_name",
|
|
)
|
|
assert updated_webhook.name == "updated_name"
|
|
assert updated_webhook.url == "https://example.com/original"
|
|
assert updated_webhook.events == events
|
|
assert updated_webhook.description == "Original description"
|
|
|
|
|
|
def test_update_webhook_not_found(client: MlflowClient):
|
|
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
|
|
client.update_webhook(webhook_id="nonexistent", name="new_name")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("invalid_url", "expected_match"),
|
|
[
|
|
(" ", r"Webhook URL cannot be empty or just whitespace"),
|
|
("ftp://example.com", r"Invalid webhook URL scheme"),
|
|
("http://[invalid", r"Invalid webhook URL"),
|
|
],
|
|
)
|
|
def test_update_webhook_invalid_urls(client: MlflowClient, invalid_url: str, expected_match: str):
|
|
# Create a valid webhook first
|
|
webhook = client.create_webhook(
|
|
name="test_webhook",
|
|
url="https://example.com/webhook",
|
|
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
|
|
)
|
|
with pytest.raises(MlflowException, match=expected_match):
|
|
client.update_webhook(webhook_id=webhook.webhook_id, url=invalid_url)
|
|
|
|
|
|
def test_delete_webhook(client: MlflowClient):
|
|
events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)]
|
|
webhook = client.create_webhook(
|
|
name="test_webhook",
|
|
url="https://example.com/webhook",
|
|
events=events,
|
|
)
|
|
client.delete_webhook(webhook.webhook_id)
|
|
with pytest.raises(MlflowException, match=r"Webhook with ID .* not found"):
|
|
client.get_webhook(webhook.webhook_id)
|
|
webhooks_page = client.list_webhooks()
|
|
webhook_ids = {w.webhook_id for w in webhooks_page}
|
|
assert webhook.webhook_id not in webhook_ids
|
|
|
|
|
|
def test_delete_webhook_not_found(client: MlflowClient):
|
|
with pytest.raises(MlflowException, match="Webhook with ID nonexistent not found"):
|
|
client.delete_webhook("nonexistent")
|