portfolio-optimiser/tests/conftest.py

"""Shared e2e fixtures (Step 13): a scripted chat client that emits a SYNTHETIC UsageDetails
(so token accounting is real-shaped without an LLM), plus store + docs-dir fixtures.

The synthetic ``UsageDetails`` is what lets the budget meter / provenance ``token_usage`` be a
positive, UsageDetails-sourced number in CI — the REAL-provider populated-usage assertion is
the gated live arm (Step 14).
"""

from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import Any

import pytest
from agent_framework import (
    BaseChatClient,
    ChatResponse,
    ChatResponseUpdate,
    Message,
    UsageDetails,
)
from agent_framework_openai import OpenAIChatCompletionClient

from portfolio_optimiser.verdicts import VerdictStore, seed_store


class SyntheticUsageChatClient(OpenAIChatCompletionClient):
    """Network-free chat client returning scripted/default replies WITH a synthetic
    ``UsageDetails`` (``total_token_count``), so strict usage accounting does not hard-fail.

    Subclasses the LAYERED ``OpenAIChatCompletionClient`` (not the minimal ``BaseChatClient``)
    so it inherits the ``ChatMiddlewareLayer`` — a ``ChatMiddleware`` attached to an agent
    backed by the minimal base would silently no-op (verified). Construction is offline
    (loopback ``base_url``, dummy key); the ``_inner_get_response`` override intercepts the
    raw call before any HTTP, so no network is touched."""

    OTEL_PROVIDER_NAME = "synthetic"

    def __init__(
        self,
        scripted: Sequence[str] | None = None,
        *,
        default_reply: str = "ok",
        tokens_per_reply: int = 8,
    ) -> None:
        super().__init__(model="synthetic", api_key="synthetic", base_url="http://127.0.0.1:9/v1")
        self._scripted = list(scripted or [])
        self._idx = 0
        self._default = default_reply
        self._tokens = tokens_per_reply
        self.call_count = 0

    def _next_reply(self) -> str:
        reply = self._scripted[self._idx] if self._idx < len(self._scripted) else self._default
        self._idx += 1
        self.call_count += 1
        return reply

    def _inner_get_response(
        self, *, messages: Sequence[Message], stream: bool, options: Any, **kwargs: Any
    ) -> Any:
        reply = self._next_reply()
        usage = UsageDetails(total_token_count=self._tokens)
        if stream:

            async def _agen() -> Any:
                yield ChatResponseUpdate(
                    role="assistant", contents=[{"type": "text", "text": reply}]
                )

            return self._build_response_stream(_agen())

        async def _coro() -> ChatResponse:
            return ChatResponse(
                messages=[Message(role="assistant", contents=[reply])],
                response_id="synthetic",
                usage_details=usage,
            )

        return _coro()


@pytest.fixture()
def make_client_factory() -> Callable[..., Callable[[str], BaseChatClient]]:
    """Return a maker that builds a per-role client factory emitting synthetic usage."""

    def _make(default_reply: str, *, tokens: int = 8) -> Callable[[str], BaseChatClient]:
        def factory(role: str) -> BaseChatClient:
            return SyntheticUsageChatClient(default_reply=default_reply, tokens_per_reply=tokens)

        return factory

    return _make


# A generic VALID SavingsProposal reply for any project not present in a portfolio reply map:
# affected total = 1 x 100_000 = 100_000, P90 = 0.30 x 100_000 = 30_000, claimed 20_000 <= both
# (Pydantic affected-total invariant and the validator P90 gate) -> always validates.
_PORTFOLIO_DEFAULT_REPLY = (
    '{"measure":"Reduce scope","affected_items":'
    '[{"code":"01.1","quantity":1,"unit_cost":100000}],"claimed_saving_nok":20000}'
)


class _ProjectAwareUsageChatClient(SyntheticUsageChatClient):
    """A ``SyntheticUsageChatClient`` that selects its reply by scanning the incoming prompt for
    a known ``project_id`` substring (the prompt embeds ``project.id`` at run.py:162 and
    generate.py:48), falling back to a default valid proposal. This keeps ``run_portfolio``'s
    single ``client_factory`` production-shaped while letting tests vary the proposal per
    project."""

    def __init__(
        self, replies: dict[str, str], *, default_reply: str, tokens_per_reply: int = 8
    ) -> None:
        super().__init__(default_reply=default_reply, tokens_per_reply=tokens_per_reply)
        self._replies = dict(replies)

    def _inner_get_response(
        self, *, messages: Sequence[Message], stream: bool, options: Any, **kwargs: Any
    ) -> Any:
        blob = " ".join(getattr(m, "text", "") or "" for m in messages)
        reply = next((r for pid, r in self._replies.items() if pid in blob), self._default)
        self.call_count += 1
        usage = UsageDetails(total_token_count=self._tokens)
        if stream:

            async def _agen() -> Any:
                yield ChatResponseUpdate(
                    role="assistant", contents=[{"type": "text", "text": reply}]
                )

            return self._build_response_stream(_agen())

        async def _coro() -> ChatResponse:
            return ChatResponse(
                messages=[Message(role="assistant", contents=[reply])],
                response_id="synthetic",
                usage_details=usage,
            )

        return _coro()


@pytest.fixture()
def make_portfolio_client_factory() -> Callable[..., Callable[[str], BaseChatClient]]:
    """Return a maker that builds a single project-aware client factory: every client it
    produces picks its reply from ``replies`` by scanning the prompt for the project id, so one
    factory serves the whole portfolio (matching ``run_portfolio``'s single-factory seam)."""

    def _make(
        replies: dict[str, str],
        *,
        default_reply: str = _PORTFOLIO_DEFAULT_REPLY,
        tokens: int = 8,
    ) -> Callable[[str], BaseChatClient]:
        def factory(role: str) -> BaseChatClient:
            return _ProjectAwareUsageChatClient(
                replies, default_reply=default_reply, tokens_per_reply=tokens
            )

        return factory

    return _make


@pytest.fixture()
def fresh_store() -> VerdictStore:
    return VerdictStore(verdicts=[])


@pytest.fixture()
def seeded_store() -> VerdictStore:
    return seed_store()


@pytest.fixture()
def docs_dir(tmp_path) -> str:
    d = tmp_path / "docs"
    d.mkdir()
    (d / "cost.txt").write_text(
        "Asphalt Ab11 unit rate renegotiation reduced the paving cost on the school stretch.",
        encoding="utf-8",
    )
    return str(d)