77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import timedelta
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
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 * * * *",
|
|
)
|
|
|
|
assert captured["schedule"].cron == "0 * * * *"
|
|
assert captured["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"],
|
|
)
|
|
|
|
assert captured["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: object, limit: int) -> None:
|
|
called["deployments"] = deployments
|
|
called["limit"] = limit
|
|
|
|
monkeypatch.setattr(scheduler, "prefect_serve", fake_serve)
|
|
|
|
deployment = SimpleNamespace(name="only")
|
|
scheduler.serve_deployments([deployment])
|
|
|
|
assert called["deployments"] == (deployment,)
|
|
assert called["limit"] == 10
|