tests/test_server_observability.py

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