atlas-iac/testing/tests/test_hermes_ai_usage_codex.py

"""Adversarial tests for bounded Codex quota-query process handling."""

from __future__ import annotations

import importlib.util
import json
import subprocess
import sys
from io import StringIO
from pathlib import Path
from types import SimpleNamespace

import pytest


ROOT = Path(__file__).resolve().parents[2]
SCRIPT = ROOT / "services/hermes/scripts/ai_usage_codex.py"


def load_module():
    """Load the standalone Codex query helper as a fresh test module."""
    name = "ai_usage_codex_tested"
    spec = importlib.util.spec_from_file_location(name, SCRIPT)
    module = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def test_default_query_deadline_has_margin_for_live_app_server_jitter(monkeypatch):
    """The default must not sit on the live app-server's common 20s boundary."""
    monkeypatch.setenv("ATLAS_AI_CODEX_QUERY_TIMEOUT_SECONDS", "45")
    mod = load_module()

    assert mod.CODEX_QUERY_TIMEOUT_SECONDS == 45


def test_codex_query_uses_structured_app_server_protocol(tmp_path, monkeypatch):
    mod = load_module()
    request_log = tmp_path / "requests.jsonl"
    mock = tmp_path / "codex-mock"
    mock.write_text(
        """#!/usr/bin/env python3
import json, os, sys
messages = [json.loads(sys.stdin.readline()) for _ in range(4)]
with open(os.environ['REQUEST_LOG'], 'w') as handle:
    for message in messages:
        handle.write(json.dumps(message) + '\\n')
responses = {
    1: {},
    2: {'rateLimits': {'primary': {'usedPercent': 10}}},
    3: {'summary': {}, 'dailyUsageBuckets': []},
}
for response_id, result in responses.items():
    print(json.dumps({'id': response_id, 'result': result}), flush=True)
"""
    )
    mock.chmod(0o755)
    monkeypatch.setattr(mod, "CODEX_BIN", str(mock))
    monkeypatch.setenv("REQUEST_LOG", str(request_log))

    rate, usage = mod.query_codex(timeout=5)
    requests = [json.loads(line) for line in request_log.read_text().splitlines()]

    assert rate["rateLimits"]["primary"]["usedPercent"] == 10
    assert usage["summary"] == {}
    assert [request["method"] for request in requests] == [
        "initialize",
        "initialized",
        "account/rateLimits/read",
        "account/usage/read",
    ]
    assert all("/status" not in json.dumps(request) for request in requests)


def test_codex_cleanup_is_bounded_when_process_cannot_be_reaped(monkeypatch):
    mod = load_module()
    processes = []
    selectors = []

    class StubbornProcess:
        def __init__(self, *_args, **_kwargs):
            self.stdin = StringIO()
            self.stdout = StringIO(
                "\n".join(
                    json.dumps({"id": response_id, "result": result})
                    for response_id, result in (
                        (1, {}),
                        (2, {"rateLimits": {}}),
                        (3, {"summary": {}, "dailyUsageBuckets": []}),
                    )
                )
                + "\n"
            )
            self.stderr = None
            self.terminate_calls = 0
            self.kill_calls = 0
            self.wait_timeouts = []
            processes.append(self)

        def terminate(self):
            self.terminate_calls += 1

        def kill(self):
            self.kill_calls += 1

        def wait(self, timeout):
            self.wait_timeouts.append(timeout)
            raise subprocess.TimeoutExpired("codex-mock", timeout)

    class TrackingSelector:
        def __init__(self):
            self.fileobj = None
            self.closed = False
            selectors.append(self)

        def register(self, fileobj, _events):
            self.fileobj = fileobj

        def select(self, _timeout):
            return [(SimpleNamespace(fileobj=self.fileobj), None)]

        def close(self):
            self.closed = True

    monkeypatch.setattr(mod.subprocess, "Popen", StubbornProcess)
    monkeypatch.setattr(mod.selectors, "DefaultSelector", TrackingSelector)

    with pytest.raises(RuntimeError, match="^Codex app-server cleanup failed$"):
        mod.query_codex(timeout=1)

    process = processes[0]
    assert process.terminate_calls == 1
    assert process.kill_calls == 1
    assert process.wait_timeouts == [5, 5]
    assert process.stdin.closed
    assert process.stdout.closed
    assert selectors[0].closed


def test_cleanup_contains_all_pipe_and_process_api_failures():
    mod = load_module()

    class BrokenStream:
        def close(self):
            raise OSError("private pipe detail")

    class BrokenSelector:
        def close(self):
            raise OSError("private selector detail")

    class BrokenProcess:
        stdin = BrokenStream()
        stdout = BrokenStream()
        stderr = BrokenStream()

        def __init__(self):
            self.wait_calls = 0

        def terminate(self):
            raise OSError("private terminate detail")

        def wait(self, timeout):
            assert timeout == mod.CODEX_CLEANUP_TIMEOUT_SECONDS
            self.wait_calls += 1
            raise OSError("private wait detail")

        def kill(self):
            raise OSError("private kill detail")

    assert mod._close_stream(None)
    assert not mod._close_stream(BrokenStream())
    assert not mod._cleanup_codex_process(BrokenProcess(), BrokenSelector())


def test_cleanup_accepts_already_exited_process_during_terminate_and_kill():
    mod = load_module()

    class AlreadyExited:
        stdin = None
        stdout = None
        stderr = None

        def __init__(self, *, timeout_first):
            self.timeout_first = timeout_first
            self.wait_calls = 0

        def terminate(self):
            raise ProcessLookupError

        def wait(self, timeout):
            self.wait_calls += 1
            if self.timeout_first and self.wait_calls == 1:
                raise subprocess.TimeoutExpired("codex", timeout)

        def kill(self):
            raise ProcessLookupError

    assert mod._cleanup_codex_process(AlreadyExited(timeout_first=False), None)
    assert mod._cleanup_codex_process(AlreadyExited(timeout_first=True), None)


def test_query_rejects_missing_pipes_and_structured_errors(monkeypatch):
    mod = load_module()

    class Process:
        stderr = None

        def __init__(self, *_args, **_kwargs):
            self.stdin = None
            self.stdout = None

        def terminate(self):
            return None

        def wait(self, timeout):
            return 0

    monkeypatch.setattr(mod.subprocess, "Popen", Process)
    with pytest.raises(RuntimeError, match="pipes are unavailable"):
        mod.query_codex(timeout=0)

    class InvalidProcess(Process):
        def __init__(self, *_args, **_kwargs):
            self.stdin = StringIO()
            self.stdout = SequencedOutput()

    class SequencedOutput:
        def __init__(self):
            self.lines = iter(
                [
                    "",
                    json.dumps({"id": 99, "result": {}}) + "\n",
                    json.dumps({"id": 1, "result": {}}) + "\n",
                    json.dumps({"id": 2, "error": {"message": "secret"}}) + "\n",
                    json.dumps({"id": 3, "result": {}}) + "\n",
                ]
            )

        def readline(self):
            return next(self.lines)

        def close(self):
            return None

    class Selector:
        def register(self, fileobj, _events):
            self.fileobj = fileobj

        def select(self, _timeout):
            return [(SimpleNamespace(fileobj=self.fileobj), None)]

        def close(self):
            return None

    monkeypatch.setattr(mod.subprocess, "Popen", InvalidProcess)
    monkeypatch.setattr(mod.selectors, "DefaultSelector", Selector)
    with pytest.raises(
        RuntimeError, match="^Codex app-server response 2 failed$"
    ) as caught:
        mod.query_codex(timeout=1)
    assert "secret" not in str(caught.value)