arc53--docsgpt/tests/storage/db/repositories/test_stack_logs.py

"""Tests for StackLogsRepository against a real Postgres instance."""

from __future__ import annotations

from sqlalchemy import text

from application.storage.db.repositories.stack_logs import StackLogsRepository


def _repo(conn) -> StackLogsRepository:
    return StackLogsRepository(conn)


class TestInsert:
    def test_inserts_log(self, pg_conn):
        repo = _repo(pg_conn)
        repo.insert(
            activity_id="act-1",
            endpoint="/api/answer",
            level="info",
            user_id="u1",
            api_key="k1",
            query="what is python?",
            stacks=[{"component": "retriever", "data": {"docs": 3}}],
        )
        row = pg_conn.execute(
            text("SELECT * FROM stack_logs WHERE activity_id = 'act-1'")
        ).fetchone()
        assert row is not None
        mapping = dict(row._mapping)
        assert mapping["endpoint"] == "/api/answer"
        assert mapping["level"] == "info"
        assert mapping["user_id"] == "u1"
        assert mapping["stacks"] == [{"component": "retriever", "data": {"docs": 3}}]

    def test_inserts_with_empty_stacks(self, pg_conn):
        repo = _repo(pg_conn)
        repo.insert(activity_id="act-2", level="error")
        row = pg_conn.execute(
            text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-2'")
        ).fetchone()
        assert row is not None
        assert dict(row._mapping)["stacks"] == []

    def test_truncated_query_stored(self, pg_conn):
        repo = _repo(pg_conn)
        long_query = "x" * 20000
        repo.insert(activity_id="act-3", query=long_query)
        row = pg_conn.execute(
            text("SELECT query FROM stack_logs WHERE activity_id = 'act-3'")
        ).fetchone()
        assert len(dict(row._mapping)["query"]) == 20000

    def test_secrets_redacted_from_stacks(self, pg_conn):
        # The ``llm`` stack component is built from every public attr of
        # the LLM and includes the provider secret + the caller's key.
        # These must never persist (the unified-logs endpoint returns
        # stacks verbatim).
        repo = _repo(pg_conn)
        repo.insert(
            activity_id="act-secret",
            level="error",
            stacks=[
                {
                    "component": "llm",
                    "data": {
                        "api_key": "sk-deployment-secret",
                        "user_api_key": "agent-key",
                        "OPENAI_API_KEY": "sk-env",
                        "model": "gpt-x",
                        "prompt_tokens": 42,
                    },
                }
            ],
        )
        row = pg_conn.execute(
            text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-secret'")
        ).fetchone()
        data = dict(row._mapping)["stacks"][0]["data"]
        assert data["api_key"] == "[REDACTED]"
        assert data["user_api_key"] == "[REDACTED]"
        assert data["OPENAI_API_KEY"] == "[REDACTED]"
        # Non-secret fields (incl. token *counts*) are untouched.
        assert data["model"] == "gpt-x"
        assert data["prompt_tokens"] == 42

    def test_redacts_broadened_secret_keys(self, pg_conn):
        # Credentials beyond ``*api_key`` — OAuth tokens, bearer/auth
        # headers, private keys — must also be scrubbed, while token
        # *count* fields are preserved.
        repo = _repo(pg_conn)
        repo.insert(
            activity_id="act-broad",
            level="error",
            stacks=[
                {
                    "component": "tool",
                    "data": {
                        "access_token": "ya29.secret",
                        "api_token": "r8_secret",
                        "token": "hf_secret",
                        "Authorization": "Bearer xyz",
                        "private_key": "-----BEGIN-----",
                        "client_secret": "cs_secret",
                        "prompt_tokens": 10,
                        "generated_tokens": 5,
                        "token_budget": 1000,
                    },
                }
            ],
        )
        row = pg_conn.execute(
            text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-broad'")
        ).fetchone()
        data = dict(row._mapping)["stacks"][0]["data"]
        for key in (
            "access_token",
            "api_token",
            "token",
            "Authorization",
            "private_key",
            "client_secret",
        ):
            assert data[key] == "[REDACTED]", key
        # Token *counts* / budgets are not credentials and must survive.
        assert data["prompt_tokens"] == 10
        assert data["generated_tokens"] == 5
        assert data["token_budget"] == 1000