mlflow--mlflow/tests/webhooks/test_delivery.py

import json
from pathlib import Path
from unittest.mock import Mock, patch

import pytest

from mlflow.entities.webhook import Webhook, WebhookAction, WebhookEntity, WebhookEvent
from mlflow.environment_variables import MLFLOW_ENABLE_WORKSPACES
from mlflow.store.model_registry.file_store import FileStore
from mlflow.store.model_registry.sqlalchemy_store import SqlAlchemyStore
from mlflow.webhooks.delivery import _send_webhook_request, deliver_webhook
from mlflow.webhooks.delivery import test_webhook as send_test_webhook


@pytest.fixture
def file_store(tmp_path: Path) -> FileStore:
    pytest.skip("FileStore is no longer supported.")
    return FileStore(str(tmp_path))


@pytest.fixture
def sql_store(tmp_path: Path) -> SqlAlchemyStore:
    db_file = tmp_path / "test.db"
    db_uri = f"sqlite:///{db_file}"
    return SqlAlchemyStore(db_uri)


@pytest.fixture
def webhook_event() -> WebhookEvent:
    return WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED)


@pytest.fixture
def webhook_payload() -> dict[str, str]:
    return {"name": "test_model", "description": "Test model"}


@pytest.fixture
def webhook(webhook_event: WebhookEvent) -> Webhook:
    return Webhook(
        webhook_id="wh-1",
        name="test",
        url="https://example.com/hook",
        events=[webhook_event],
        creation_timestamp=0,
        last_updated_timestamp=0,
        workspace="team-a",
    )


def test_deliver_webhook_exits_early_for_file_store(
    file_store: FileStore, webhook_event: WebhookEvent, webhook_payload: dict[str, str]
) -> None:
    pytest.skip("FileStore is no longer supported.")
    with patch("mlflow.webhooks.delivery._deliver_webhook_impl") as mock_impl:
        deliver_webhook(
            event=webhook_event,
            payload=webhook_payload,
            store=file_store,
        )

        # _deliver_webhook_impl should not be called for FileStore
        mock_impl.assert_not_called()


def test_deliver_webhook_calls_impl_for_sql_store(
    sql_store: SqlAlchemyStore, webhook_event: WebhookEvent, webhook_payload: dict[str, str]
) -> None:
    with patch("mlflow.webhooks.delivery._deliver_webhook_impl") as mock_impl:
        deliver_webhook(
            event=webhook_event,
            payload=webhook_payload,
            store=sql_store,
        )

        # _deliver_webhook_impl should be called for SqlAlchemyStore
        mock_impl.assert_called_once_with(
            event=webhook_event,
            payload=webhook_payload,
            store=sql_store,
        )


def test_deliver_webhook_handles_exception_for_sql_store(
    sql_store: SqlAlchemyStore, webhook_event: WebhookEvent, webhook_payload: dict[str, str]
) -> None:
    with (
        patch("mlflow.webhooks.delivery._deliver_webhook_impl", side_effect=Exception("Test")),
        patch("mlflow.webhooks.delivery._logger") as mock_logger,
    ):
        # This should not raise an exception
        deliver_webhook(
            event=webhook_event,
            payload=webhook_payload,
            store=sql_store,
        )

        # Verify that the error was logged
        mock_logger.error.assert_called_once()
        assert "Failed to deliver webhook for event" in str(mock_logger.error.call_args)


def test_deliver_webhook_no_exception_for_file_store(
    file_store: FileStore, webhook_event: WebhookEvent, webhook_payload: dict[str, str]
) -> None:
    pytest.skip("FileStore is no longer supported.")
    with (
        patch(
            "mlflow.webhooks.delivery._deliver_webhook_impl", side_effect=Exception("Test")
        ) as mock_impl,
        patch("mlflow.webhooks.delivery._logger") as mock_logger,
    ):
        # This should not raise an exception and should return early
        deliver_webhook(
            event=webhook_event,
            payload=webhook_payload,
            store=file_store,
        )

        # _deliver_webhook_impl should not be called, so no error should be logged
        mock_impl.assert_not_called()
        mock_logger.error.assert_not_called()


def test_test_webhook_rejects_private_ip():
    event = WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)
    webhook = Webhook(
        webhook_id="wh-1",
        name="test",
        url="https://localhost/hook",
        events=[event],
        creation_timestamp=0,
        last_updated_timestamp=0,
    )

    result = send_test_webhook(webhook)

    assert result.success is False
    assert "must not resolve to a non-public" in result.error_message


def test_send_webhook_request_includes_workspace_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
    webhook: Webhook,
    webhook_event: WebhookEvent,
    webhook_payload: dict[str, str],
) -> None:
    monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
    session = Mock()
    session.post.return_value = Mock(status_code=200, text="ok")

    with patch("mlflow.webhooks.delivery._validate_webhook_url"):
        _send_webhook_request(webhook, webhook_payload, webhook_event, session)

    actual_payload = json.loads(session.post.call_args.kwargs["data"].decode("utf-8"))
    assert actual_payload["entity"] == webhook_event.entity.value
    assert actual_payload["action"] == webhook_event.action.value
    assert actual_payload["workspace"] == webhook.workspace
    assert actual_payload["data"] == webhook_payload
    assert "timestamp" in actual_payload


def test_send_webhook_request_omits_workspace_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
    webhook: Webhook,
    webhook_event: WebhookEvent,
    webhook_payload: dict[str, str],
) -> None:
    monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "false")
    session = Mock()
    session.post.return_value = Mock(status_code=200, text="ok")

    with patch("mlflow.webhooks.delivery._validate_webhook_url"):
        _send_webhook_request(webhook, webhook_payload, webhook_event, session)

    actual_payload = json.loads(session.post.call_args.kwargs["data"].decode("utf-8"))
    assert "workspace" not in actual_payload