89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import timedelta
|
|
from types import SimpleNamespace
|
|
from typing import cast
|
|
|
|
import pytest
|
|
from prefect.deployments.runner import RunnerDeployment
|
|
from prefect.schedules import Cron, Interval
|
|
|
|
from ingest_pipeline.flows import scheduler
|
|
|
|
|
|
def test_create_scheduled_deployment_cron(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured: dict[str, object] = {}
|
|
|
|
class DummyFlow:
|
|
def to_deployment(self, **kwargs: object) -> SimpleNamespace:
|
|
nonlocal captured
|
|
captured |= kwargs
|
|
return SimpleNamespace(**kwargs)
|
|
|
|
monkeypatch.setattr(scheduler, "create_ingestion_flow", DummyFlow())
|
|
|
|
deployment = scheduler.create_scheduled_deployment(
|
|
name="cron-ingestion",
|
|
source_url="https://example.com",
|
|
source_type="web",
|
|
schedule_type="cron",
|
|
cron_expression="0 * * * *",
|
|
)
|
|
|
|
schedule = captured["schedule"]
|
|
# Check that it's a cron schedule by verifying it has a cron attribute
|
|
assert hasattr(schedule, "cron")
|
|
assert schedule.cron == "0 * * * *"
|
|
|
|
parameters = captured["parameters"]
|
|
assert isinstance(parameters, dict)
|
|
assert parameters["source_type"] == "web"
|
|
assert deployment.tags == ["web", "weaviate"]
|
|
|
|
|
|
def test_create_scheduled_deployment_interval(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured: dict[str, object] = {}
|
|
|
|
class DummyFlow:
|
|
def to_deployment(self, **kwargs: object) -> SimpleNamespace:
|
|
nonlocal captured
|
|
captured |= kwargs
|
|
return SimpleNamespace(**kwargs)
|
|
|
|
monkeypatch.setattr(scheduler, "create_ingestion_flow", DummyFlow())
|
|
|
|
deployment = scheduler.create_scheduled_deployment(
|
|
name="interval-ingestion",
|
|
source_url="https://repo.example.com",
|
|
source_type="repository",
|
|
storage_backend="open_webui",
|
|
schedule_type="interval",
|
|
interval_minutes=15,
|
|
tags=["custom"],
|
|
)
|
|
|
|
schedule = captured["schedule"]
|
|
# Check that it's an interval schedule by verifying it has an interval attribute
|
|
assert hasattr(schedule, "interval")
|
|
assert schedule.interval == timedelta(minutes=15)
|
|
|
|
assert captured["tags"] == ["custom"]
|
|
assert deployment.parameters["storage_backend"] == "open_webui"
|
|
|
|
|
|
def test_serve_deployments_invokes_prefect(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
called: dict[str, object] = {}
|
|
|
|
def fake_serve(*deployments: RunnerDeployment, limit: int) -> None:
|
|
called["deployments"] = deployments
|
|
called["limit"] = limit
|
|
|
|
monkeypatch.setattr(scheduler, "prefect_serve", fake_serve)
|
|
|
|
# Create a mock deployment using SimpleNamespace to avoid Prefect complexity
|
|
deployment = SimpleNamespace(name="only")
|
|
scheduler.serve_deployments([cast(RunnerDeployment, deployment)])
|
|
|
|
assert called["deployments"] == (deployment,)
|
|
assert called["limit"] == 10
|