from __future__ import annotations import contextlib from fastapi.testclient import TestClient from ergo_agentic import server from ergo_agentic.domain.outcome import OutcomeMatrix class _FakeGraph: def __init__(self) -> None: self.calls: list[dict] = [] async def ainvoke(self, initial_state, config): self.calls.append({"initial_state": initial_state, "config": config}) return { "outcome_matrix": OutcomeMatrix.model_validate( { "scene_config": { "screen_count": 0, "screen_types": [], "has_standing_desk": False, "person_detected": False, }, "setup_summary": "", "assessments": [], "skipped_parameters": [ { "parameter_id": "monitor-height", "parameter_text": "Monitor height", "reason": "no screen visible", } ], } ) } class _FakeActiveTrace: def __init__(self) -> None: self.completed = None self.failed = None self.graph_entered = 0 @contextlib.contextmanager def graph_invoke(self): self.graph_entered += 1 yield def complete(self, matrix): self.completed = matrix def fail(self, exc): self.failed = exc class _FakeObserver: def __init__(self) -> None: self.calls: list[dict] = [] self.active = _FakeActiveTrace() @contextlib.contextmanager def observe_assessment(self, *, run_id, images, metadata): self.calls.append({"run_id": run_id, "images": images, "metadata": metadata}) yield self.active def test_assess_runs_through_observability_wrapper(): fake_graph = _FakeGraph() fake_observer = _FakeObserver() with TestClient(server.app) as client: previous_graph = server._graph previous_observer = server._workflow_observer server._graph = fake_graph server._workflow_observer = fake_observer try: res = client.post( "/assess", json={ "run_id": "run_poc", "images": [{"image_id": "img_1", "url": "data:image/png;base64,SECRET"}], "metadata": { "user_name": "Alice Secret", "equipment": {"screen_count": 1, "has_standing_desk": False}, }, }, ) finally: server._graph = previous_graph server._workflow_observer = previous_observer assert res.status_code == 200, res.text assert fake_observer.calls[0]["run_id"] == "run_poc" assert fake_observer.calls[0]["metadata"]["user_name"] == "Alice Secret" assert fake_observer.active.graph_entered == 1 assert fake_observer.active.completed is not None assert fake_observer.active.failed is None