pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/testing/mocks.py

@@ -0,0 +1,104 @@
+"""Mock implementations of engine ABCs for product testing."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pulse_engine.extractor.base import BaseExtractor, ExtractionResult, ExtractorConfig
+from pulse_engine.extractor.orchestrator.base import (
+    BaseOrchestratorAdapter,
+    OrchestratorRunStatus,
+)
+from pulse_engine.storage.connectors.base import BaseStorageConnector
+from pulse_engine.storage.schemas import (
+    ConnectorHealth,
+    Document,
+    SearchQuery,
+    SearchResult,
+    StoreResult,
+)
+
+
+class MockStorageConnector(BaseStorageConnector):
+    """In-memory storage connector for testing."""
+
+    def __init__(self) -> None:
+        self._documents: dict[str, dict[str, Document]] = {}
+
+    async def initialize(self) -> None:
+        pass
+
+    async def store(self, tenant_id: str, documents: list[Document]) -> StoreResult:
+        bucket = self._documents.setdefault(tenant_id, {})
+        stored = 0
+        for doc in documents:
+            key = f"{doc.doc_id}:{doc.chunk_id}" if doc.chunk_id else doc.doc_id
+            bucket[key] = doc
+            stored += 1
+        return StoreResult(stored_count=stored, failed_count=0)
+
+    async def retrieve(self, tenant_id: str, doc_id: str) -> Document | None:
+        bucket = self._documents.get(tenant_id, {})
+        return bucket.get(doc_id)
+
+    async def search(self, tenant_id: str, query: SearchQuery) -> SearchResult:
+        return SearchResult(hits=[], total=0, took_ms=0)
+
+    async def delete(self, tenant_id: str, doc_id: str) -> bool:
+        bucket = self._documents.get(tenant_id, {})
+        return bucket.pop(doc_id, None) is not None
+
+    async def health_check(self) -> ConnectorHealth:
+        return ConnectorHealth(connector="mock", status="up", latency_ms=0.0)
+
+    async def teardown(self) -> None:
+        self._documents.clear()
+
+
+class MockOrchestratorAdapter(BaseOrchestratorAdapter):
+    """No-op orchestrator adapter for testing."""
+
+    def __init__(self) -> None:
+        self._runs: dict[str, str] = {}
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        status = self._runs.get(run_id, "unknown")
+        return OrchestratorRunStatus(run_id=run_id, status=status)
+
+    async def cancel_run(self, run_id: str) -> bool:
+        self._runs[run_id] = "cancelled"
+        return True
+
+    async def health_check(self) -> bool:
+        return True
+
+    async def create_flow_run(
+        self,
+        deployment_id: str,
+        parameters: dict[str, Any] | None = None,
+    ) -> str:
+        return ""
+
+
+class MockExtractor(BaseExtractor):
+    """A configurable mock extractor for testing."""
+
+    def __init__(
+        self,
+        name: str = "mock",
+        job_type: str = "mock_extract",
+        results: list[ExtractionResult] | None = None,
+    ) -> None:
+        self._name = name
+        self._job_type = job_type
+        self._results = results or []
+        self.extract_calls: list[tuple[str, dict[str, Any]]] = []
+
+    def get_config(self) -> ExtractorConfig:
+        return ExtractorConfig(name=self._name, job_type=self._job_type)
+
+    async def extract(
+        self, tenant_id: str, parameters: dict[str, Any]
+    ) -> list[ExtractionResult]:
+        self.extract_calls.append((tenant_id, parameters))
+        return self._results

pulse_engine/worker.py

@@ -0,0 +1,53 @@
+"""Celery application factory.
+
+The engine provides the Celery infrastructure (broker, result backend, base
+configuration).  Individual *products* register their own tasks by importing
+this app and using ``@celery_app.task``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from celery import Celery
+
+if TYPE_CHECKING:
+    from pulse_engine.config import Settings
+    from pulse_engine.registry import ProductManifest
+
+
+def create_celery_app(
+    settings: Settings, manifest: ProductManifest | None = None
+) -> Celery:
+    app = Celery(
+        "pulse",
+        broker=settings.effective_celery_broker_url,
+        backend=settings.effective_celery_result_backend,
+    )
+
+    app.conf.update(
+        task_serializer="json",
+        accept_content=["json"],
+        result_serializer="json",
+        timezone="UTC",
+        enable_utc=True,
+        task_track_started=True,
+        task_acks_late=True,
+        worker_prefetch_multiplier=1,
+        broker_connection_retry_on_startup=True,
+    )
+
+    # Auto-discover tasks in engine and any product modules
+    packages = ["pulse_engine"]
+    if manifest and manifest.celery_task_modules:
+        packages.extend(manifest.celery_task_modules)
+    app.autodiscover_tasks(packages)
+
+    return app
+
+
+# Module-level instance for Celery CLI (-A pulse_engine.worker) and for product
+# packages to import as ``from pulse_engine.worker import celery_app``.
+from pulse_engine.config import get_settings  # noqa: E402
+
+celery_app = create_celery_app(get_settings())