151 lines
4.3 KiB
Python
151 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from types import SimpleNamespace
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from ingest_pipeline.cli import main
|
|
from ingest_pipeline.core.models import (
|
|
IngestionResult,
|
|
IngestionSource,
|
|
IngestionStatus,
|
|
StorageBackend,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("collection_arg", "expected"),
|
|
(
|
|
(None, "docs_example_com_web"),
|
|
("explicit_collection", "explicit_collection"),
|
|
),
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_run_ingestion_collection_resolution(
|
|
monkeypatch: pytest.MonkeyPatch, collection_arg: str | None, expected: str
|
|
) -> None:
|
|
recorded: dict[str, object] = {}
|
|
|
|
async def fake_flow(**kwargs: object) -> IngestionResult:
|
|
recorded.update(kwargs)
|
|
return IngestionResult(
|
|
job_id=uuid4(),
|
|
status=IngestionStatus.COMPLETED,
|
|
documents_processed=7,
|
|
documents_failed=0,
|
|
duration_seconds=1.5,
|
|
error_messages=[],
|
|
)
|
|
|
|
monkeypatch.setattr(main, "create_ingestion_flow", fake_flow)
|
|
|
|
result = await main.run_ingestion(
|
|
url="https://docs.example.com/guide",
|
|
source_type=IngestionSource.WEB,
|
|
storage_backend=StorageBackend.WEAVIATE,
|
|
collection_name=collection_arg,
|
|
validate_first=False,
|
|
)
|
|
|
|
assert recorded["collection_name"] == expected
|
|
assert result.documents_processed == 7
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("cron_value", "serve_now", "expected_type", "serve_expected"),
|
|
(
|
|
("0 * * * *", False, "cron", False),
|
|
(None, True, "interval", True),
|
|
),
|
|
)
|
|
def test_schedule_creates_deployment(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
cron_value: str | None,
|
|
serve_now: bool,
|
|
expected_type: str,
|
|
serve_expected: bool,
|
|
) -> None:
|
|
recorded: dict[str, object] = {}
|
|
served = {"called": False}
|
|
|
|
def fake_create_scheduled_deployment(**kwargs: object) -> str:
|
|
recorded.update(kwargs)
|
|
return "deployment"
|
|
|
|
def fake_serve_deployments(deployments: list[str]) -> None:
|
|
served["called"] = True
|
|
assert deployments == ["deployment"]
|
|
|
|
monkeypatch.setattr(main, "create_scheduled_deployment", fake_create_scheduled_deployment)
|
|
monkeypatch.setattr(main, "serve_deployments", fake_serve_deployments)
|
|
|
|
main.schedule(
|
|
name="nightly",
|
|
source_url="https://example.com",
|
|
source_type=IngestionSource.WEB,
|
|
storage=StorageBackend.WEAVIATE,
|
|
cron=cron_value,
|
|
interval=30,
|
|
serve_now=serve_now,
|
|
)
|
|
|
|
assert recorded["schedule_type"] == expected_type
|
|
assert served["called"] is serve_expected
|
|
|
|
|
|
def test_serve_tui_launch(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
invoked = {"count": 0}
|
|
|
|
def fake_dashboard() -> None:
|
|
invoked["count"] = invoked["count"] + 1
|
|
|
|
monkeypatch.setattr("ingest_pipeline.cli.tui.dashboard", fake_dashboard)
|
|
|
|
main.serve(ui="tui")
|
|
|
|
assert invoked["count"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_search_collects_results(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
messages: list[object] = []
|
|
|
|
class DummyConsole:
|
|
def print(self, message: object) -> None:
|
|
messages.append(message)
|
|
|
|
class WeaviateStub:
|
|
def __init__(self, config: object) -> None:
|
|
self.config = config
|
|
self.initialized = False
|
|
|
|
async def initialize(self) -> None:
|
|
self.initialized = True
|
|
|
|
async def search(
|
|
self,
|
|
query: str,
|
|
limit: int = 10,
|
|
threshold: float = 0.7,
|
|
*,
|
|
collection_name: str | None = None,
|
|
) -> AsyncGenerator[SimpleNamespace, None]:
|
|
yield SimpleNamespace(title="Title", content="Body text", score=0.91)
|
|
|
|
dummy_settings = SimpleNamespace(
|
|
weaviate_endpoint="http://weaviate.local",
|
|
weaviate_api_key="token",
|
|
openwebui_endpoint="http://chat.local",
|
|
openwebui_api_key=None,
|
|
)
|
|
|
|
monkeypatch.setattr(main, "console", DummyConsole())
|
|
monkeypatch.setattr(main, "get_settings", lambda: dummy_settings)
|
|
monkeypatch.setattr("ingest_pipeline.storage.weaviate.WeaviateStorage", WeaviateStub)
|
|
|
|
await main.run_search("query", collection=None, backend="weaviate", limit=1)
|
|
|
|
assert messages[-1] == "\n✅ [green]Found 1 results[/green]"
|