131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
"""Shared pytest fixtures for asf-worker tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from uuid import UUID, uuid4
|
|
|
|
import pytest
|
|
|
|
WORKER_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(WORKER_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(WORKER_ROOT))
|
|
|
|
TENANT_ID = UUID("00000000-0000-0000-0000-000000000099")
|
|
PROJECT_ID = UUID("00000000-0000-0000-0000-000000000088")
|
|
CORRELATION_ID = "test-correlation-worker"
|
|
|
|
|
|
class FakeJobsCollection:
|
|
"""In-memory AsyncIOMotorCollection stand-in for claim/persistence tests.
|
|
|
|
Mirrors the tiny subset of the Motor API that `_claim_running` and
|
|
`_mark_job_done` exercise: ``update_one(filter, ops)`` with
|
|
``matched_count`` on the returned UpdateResult.
|
|
"""
|
|
|
|
@dataclass
|
|
class UpdateResult:
|
|
matched_count: int = 0
|
|
modified_count: int = 0
|
|
upserted_id: object = None
|
|
|
|
def __init__(self, docs: dict[UUID, dict] | None = None) -> None:
|
|
self._docs: dict[UUID, dict] = {k: dict(v) for k, v in (docs or {}).items()}
|
|
self.last_update_filter: object = None
|
|
self.last_update_ops: object = None
|
|
|
|
async def update_one(self, filt: dict, ops: dict) -> UpdateResult:
|
|
self.last_update_filter = filt
|
|
self.last_update_ops = ops
|
|
jid = filt.get("_id")
|
|
existing = self._docs.get(jid) if jid is not None else None
|
|
if existing is None:
|
|
return FakeJobsCollection.UpdateResult(matched_count=0)
|
|
|
|
# Apply the whole filter object like MongoDB does (simple top-level equality match).
|
|
for k, v in filt.items():
|
|
if k == "status" and isinstance(v, dict) and "$in" in v:
|
|
if existing.get(k) not in v["$in"]:
|
|
return FakeJobsCollection.UpdateResult(matched_count=0)
|
|
elif k == "_id":
|
|
if existing.get(k) != v:
|
|
return FakeJobsCollection.UpdateResult(matched_count=0)
|
|
else:
|
|
if existing.get(k) != v:
|
|
return FakeJobsCollection.UpdateResult(matched_count=0)
|
|
|
|
if "$set" in ops:
|
|
for k, v in ops["$set"].items():
|
|
existing[k] = v
|
|
if "$inc" in ops:
|
|
for k, v in ops["$inc"].items():
|
|
existing[k] = (existing.get(k) or 0) + v
|
|
return FakeJobsCollection.UpdateResult(matched_count=1, modified_count=1)
|
|
|
|
|
|
class FakeDb:
|
|
def __init__(self, jobs: FakeJobsCollection | None = None) -> None:
|
|
self._jobs = jobs or FakeJobsCollection()
|
|
|
|
def __getitem__(self, col: str):
|
|
if col == "jobs":
|
|
return self._jobs
|
|
raise KeyError(col)
|
|
|
|
|
|
class FakeSettings:
|
|
mongo_uri = "mongodb://fake"
|
|
mongo_db = "fake_asf"
|
|
rabbitmq_uri = "amqp://fake"
|
|
queues: list[str] = []
|
|
prefetch_count: int = 1
|
|
worker_id: str = "test-worker"
|
|
enable_ai_handlers: bool = False
|
|
ai_gateway_provider: str = "mock"
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_settings():
|
|
return FakeSettings()
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_jobs_collection():
|
|
return FakeJobsCollection()
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_db(fake_jobs_collection):
|
|
return FakeDb(jobs=fake_jobs_collection)
|
|
|
|
|
|
@pytest.fixture
|
|
def worker_deps(fake_settings, fake_db, fake_jobs_collection):
|
|
from app.handlers import WorkerDependencies
|
|
return WorkerDependencies(
|
|
db=fake_db,
|
|
jobs=fake_jobs_collection,
|
|
ai_executions=fake_db["jobs"], # fine - not exercised
|
|
artifacts=fake_db["jobs"],
|
|
settings=fake_settings,
|
|
worker_id=fake_settings.worker_id,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_envelope():
|
|
from asf_events.base import JobRequestEnvelope
|
|
from asf_domain import JobType
|
|
return JobRequestEnvelope(
|
|
job_id=uuid4(),
|
|
tenant_id=TENANT_ID,
|
|
project_id=PROJECT_ID,
|
|
correlation_id=CORRELATION_ID,
|
|
type=JobType.DEVOPS.value,
|
|
payload={"hello": "world", "sequence": 7},
|
|
priority="normal",
|
|
)
|