mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
7fb6d7fefd
Codex treats env_key as a hard requirement: if the named env var is absent it throws "Missing environment variable" before the session starts. Subscription (ChatGPT Plus) users don't have OPENAI_API_KEY set, so injecting env_key = "OPENAI_API_KEY" blocks them at startup. With env_key absent, api_key() returns Ok(None) and Codex falls through to the existing CodexAuth (OAuth for subscription, ApiKey for PAYG) — both modes authenticate correctly without startup errors.
437 lines
18 KiB
Python
437 lines
18 KiB
Python
"""Tests for `headroom wrap codex` and `headroom unwrap codex`.
|
|
|
|
These exercise the Codex-specific ``config.toml`` injection and restoration
|
|
helpers that route Codex through the Headroom proxy. They are deliberately
|
|
end-to-end-ish: the unit tests call the helpers directly against a temp
|
|
``$HOME``, and the integration tests invoke the real Click commands the same
|
|
way a user would from the shell.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
|
|
from headroom.cli import wrap as wrap_mod
|
|
from headroom.cli.main import main
|
|
|
|
|
|
def _set_test_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
home = str(tmp_path)
|
|
monkeypatch.setenv("HOME", home)
|
|
monkeypatch.setenv("USERPROFILE", home)
|
|
|
|
|
|
@pytest.fixture
|
|
def runner() -> CliRunner:
|
|
return CliRunner()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Unit tests: helpers operating on ~/.codex/config.toml
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestStripCodexHeadroomBlocks:
|
|
"""Tests for the regex-based cleanup helper."""
|
|
|
|
def test_empty_content_returns_empty(self) -> None:
|
|
assert wrap_mod._strip_codex_headroom_blocks("") == ""
|
|
|
|
def test_returns_content_unchanged_when_no_markers(self) -> None:
|
|
original = '[profiles.default]\nmodel = "gpt-4o"\n'
|
|
cleaned = wrap_mod._strip_codex_headroom_blocks(original)
|
|
# Trailing whitespace normalization only — semantic content preserved.
|
|
assert 'model = "gpt-4o"' in cleaned
|
|
assert "[profiles.default]" in cleaned
|
|
|
|
def test_removes_complete_headroom_block(self) -> None:
|
|
wrapped = (
|
|
f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
|
'model_provider = "headroom"\n'
|
|
"\n"
|
|
"[model_providers.headroom]\n"
|
|
'base_url = "http://127.0.0.1:8787/v1"\n'
|
|
f"{wrap_mod._CODEX_END_MARKER}\n"
|
|
)
|
|
assert wrap_mod._strip_codex_headroom_blocks(wrapped) == ""
|
|
|
|
def test_preserves_user_content_around_block(self) -> None:
|
|
user_pre = '[profiles.default]\nmodel = "gpt-4o"\n'
|
|
user_post = '[mcp_servers.foo]\ncommand = "echo"\n'
|
|
wrapped = (
|
|
f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
|
'model_provider = "headroom"\n'
|
|
f"{wrap_mod._CODEX_END_MARKER}\n" + user_pre + "\n"
|
|
f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
|
"[model_providers.headroom]\n"
|
|
'base_url = "http://127.0.0.1:8787/v1"\n'
|
|
f"{wrap_mod._CODEX_END_MARKER}\n" + user_post
|
|
)
|
|
cleaned = wrap_mod._strip_codex_headroom_blocks(wrapped)
|
|
assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in cleaned
|
|
assert wrap_mod._CODEX_END_MARKER not in cleaned
|
|
assert 'model = "gpt-4o"' in cleaned
|
|
assert "[mcp_servers.foo]" in cleaned
|
|
|
|
def test_removes_stray_top_level_model_provider_line(self) -> None:
|
|
# Old wrap versions left `model_provider = "headroom"` outside markers.
|
|
content = 'foo = 1\nmodel_provider = "headroom"\nbar = 2\n'
|
|
cleaned = wrap_mod._strip_codex_headroom_blocks(content)
|
|
assert 'model_provider = "headroom"' not in cleaned
|
|
assert "foo = 1" in cleaned
|
|
assert "bar = 2" in cleaned
|
|
|
|
|
|
class TestSnapshotCodexConfig:
|
|
"""Tests for ``_snapshot_codex_config_if_unwrapped``."""
|
|
|
|
def test_creates_backup_on_first_call(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "config.toml"
|
|
backup_file = tmp_path / "config.toml.headroom-backup"
|
|
config_file.write_text('model = "gpt-4o"\n')
|
|
|
|
wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
assert backup_file.exists()
|
|
assert backup_file.read_text() == 'model = "gpt-4o"\n'
|
|
|
|
def test_does_not_overwrite_existing_backup(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "config.toml"
|
|
backup_file = tmp_path / "config.toml.headroom-backup"
|
|
config_file.write_text("second-wrap content\n")
|
|
backup_file.write_text("original-pre-wrap content\n")
|
|
|
|
wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
# Backup must still contain the *original* pre-wrap content.
|
|
assert backup_file.read_text() == "original-pre-wrap content\n"
|
|
|
|
def test_no_backup_when_config_missing(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "config.toml"
|
|
backup_file = tmp_path / "config.toml.headroom-backup"
|
|
|
|
wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
assert not backup_file.exists()
|
|
|
|
def test_no_backup_when_config_already_wrapped(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "config.toml"
|
|
backup_file = tmp_path / "config.toml.headroom-backup"
|
|
config_file.write_text(
|
|
f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
|
'model_provider = "headroom"\n'
|
|
f"{wrap_mod._CODEX_END_MARKER}\n"
|
|
)
|
|
|
|
wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
# Pre-wrap snapshot must never snapshot an already-wrapped file.
|
|
assert not backup_file.exists()
|
|
|
|
|
|
class TestInjectAndRestoreRoundTrip:
|
|
"""End-to-end wrap → unwrap cycle operating directly on a temp $HOME."""
|
|
|
|
def test_wrap_unwrap_restores_empty_state(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_file = tmp_path / ".codex" / "config.toml"
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
assert config_file.exists()
|
|
assert 'model_provider = "headroom"' in config_file.read_text()
|
|
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
# No prior config existed → the injected file is fully removed.
|
|
assert status == "removed"
|
|
assert not config_file.exists()
|
|
assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists()
|
|
|
|
def test_wrap_unwrap_restores_prior_model_provider(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_dir = tmp_path / ".codex"
|
|
config_dir.mkdir()
|
|
config_file = config_dir / "config.toml"
|
|
original = (
|
|
'model_provider = "openai"\n'
|
|
"\n"
|
|
"[model_providers.openai]\n"
|
|
'name = "OpenAI"\n'
|
|
'base_url = "https://api.openai.com/v1"\n'
|
|
)
|
|
config_file.write_text(original)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
wrapped = config_file.read_text()
|
|
assert 'model_provider = "headroom"' in wrapped
|
|
assert "[model_providers.headroom]" in wrapped
|
|
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
assert status == "restored"
|
|
assert config_file.read_text() == original
|
|
assert not (config_dir / "config.toml.headroom-backup").exists()
|
|
|
|
def test_wrap_is_idempotent(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_dir = tmp_path / ".codex"
|
|
config_dir.mkdir()
|
|
config_file = config_dir / "config.toml"
|
|
original = '[profiles.default]\nmodel = "gpt-4o"\n'
|
|
config_file.write_text(original)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
wrap_mod._inject_codex_provider_config(9999) # port change
|
|
|
|
content = config_file.read_text()
|
|
# Exactly two Headroom blocks — a top-level-key block and the
|
|
# provider-table block. Re-wrapping must not duplicate them.
|
|
assert content.count(wrap_mod._CODEX_TOP_LEVEL_MARKER) == 2
|
|
assert content.count(wrap_mod._CODEX_END_MARKER) == 2
|
|
# Latest port is honoured in both keys.
|
|
assert 'base_url = "http://127.0.0.1:9999/v1"' in content
|
|
assert 'openai_base_url = "http://127.0.0.1:9999/v1"' in content
|
|
assert 'base_url = "http://127.0.0.1:8787/v1"' not in content
|
|
assert 'openai_base_url = "http://127.0.0.1:8787/v1"' not in content
|
|
# User's original content is preserved.
|
|
assert 'model = "gpt-4o"' in content
|
|
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
assert status == "restored"
|
|
assert config_file.read_text() == original
|
|
|
|
def test_unwrap_is_noop_when_never_wrapped(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
assert status == "noop"
|
|
|
|
def test_unwrap_cleans_block_without_backup(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""Handles crash-case where wrap injected but backup was wiped."""
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_dir = tmp_path / ".codex"
|
|
config_dir.mkdir()
|
|
config_file = config_dir / "config.toml"
|
|
user_content = '[profiles.default]\nmodel = "gpt-4o"\n'
|
|
config_file.write_text(
|
|
user_content + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
|
'model_provider = "headroom"\n\n'
|
|
"[model_providers.headroom]\n"
|
|
'base_url = "http://127.0.0.1:8787/v1"\n'
|
|
f"{wrap_mod._CODEX_END_MARKER}\n"
|
|
)
|
|
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
assert status == "cleaned"
|
|
cleaned = config_file.read_text()
|
|
assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in cleaned
|
|
assert wrap_mod._CODEX_END_MARKER not in cleaned
|
|
assert 'model_provider = "headroom"' not in cleaned
|
|
assert 'model = "gpt-4o"' in cleaned
|
|
|
|
def test_unwrap_handles_malformed_prior_config(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""Unwrap preserves backup content verbatim — TOML validity isn't required."""
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_dir = tmp_path / ".codex"
|
|
config_dir.mkdir()
|
|
config_file = config_dir / "config.toml"
|
|
malformed = 'this is not valid toml ][ "" \x00\n'
|
|
config_file.write_text(malformed)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
status, _ = wrap_mod._restore_codex_provider_config()
|
|
|
|
assert status == "restored"
|
|
assert config_file.read_text() == malformed
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Subscription routing: openai_base_url intercepts ChatGPT plan traffic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSubscriptionRouting:
|
|
"""Codex subscription (ChatGPT plan) bypasses OPENAI_BASE_URL and the
|
|
custom model_provider; it uses the built-in ``openai`` provider whose
|
|
base_url defaults to ``https://chatgpt.com/backend-api/codex``.
|
|
Setting ``openai_base_url`` overrides that default for all auth modes."""
|
|
|
|
def test_inject_writes_openai_base_url(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
|
|
content = (tmp_path / ".codex" / "config.toml").read_text()
|
|
assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content
|
|
|
|
def test_openai_base_url_port_updates_on_rewrap(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
wrap_mod._inject_codex_provider_config(9999)
|
|
|
|
content = (tmp_path / ".codex" / "config.toml").read_text()
|
|
assert 'openai_base_url = "http://127.0.0.1:9999/v1"' in content
|
|
assert 'openai_base_url = "http://127.0.0.1:8787/v1"' not in content
|
|
|
|
def test_openai_base_url_removed_on_unwrap(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_dir = tmp_path / ".codex"
|
|
config_dir.mkdir()
|
|
config_file = config_dir / "config.toml"
|
|
original = '[profiles.default]\nmodel = "gpt-4o"\n'
|
|
config_file.write_text(original)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in config_file.read_text()
|
|
|
|
wrap_mod._restore_codex_provider_config()
|
|
assert config_file.read_text() == original
|
|
|
|
def test_strip_cleans_orphaned_openai_base_url(self) -> None:
|
|
"""Safety net: orphaned openai_base_url lines are cleaned up."""
|
|
content = (
|
|
'[profiles.default]\nmodel = "gpt-4o"\nopenai_base_url = "http://127.0.0.1:8787/v1"\n'
|
|
)
|
|
cleaned = wrap_mod._strip_codex_headroom_blocks(content)
|
|
assert "openai_base_url" not in cleaned
|
|
assert 'model = "gpt-4o"' in cleaned
|
|
|
|
def test_no_env_key_in_injected_provider(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""env_key must be absent so Codex doesn't require OPENAI_API_KEY.
|
|
|
|
Codex treats env_key as a hard requirement — if the env var is missing
|
|
it throws "Missing environment variable" at startup. Subscription
|
|
(ChatGPT Plus) users don't have OPENAI_API_KEY set, so injecting
|
|
env_key breaks them (issue #393).
|
|
"""
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
wrap_mod._inject_codex_provider_config(8787)
|
|
|
|
content = (tmp_path / ".codex" / "config.toml").read_text()
|
|
assert "env_key" not in content
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration tests: full `headroom wrap codex` / `headroom unwrap codex`
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_wrap_codex_prepare_only_creates_backup_and_config(
|
|
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_file = tmp_path / ".codex" / "config.toml"
|
|
config_file.parent.mkdir(parents=True)
|
|
original = 'model_provider = "openai"\n'
|
|
config_file.write_text(original)
|
|
|
|
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
|
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert 'model_provider = "headroom"' in config_file.read_text()
|
|
backup = tmp_path / ".codex" / "config.toml.headroom-backup"
|
|
assert backup.exists()
|
|
assert backup.read_text() == original
|
|
|
|
|
|
def test_unwrap_codex_restores_prior_config_end_to_end(
|
|
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""The bug report, reproduced: wrap → unwrap must round-trip cleanly."""
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_file = tmp_path / ".codex" / "config.toml"
|
|
config_file.parent.mkdir(parents=True)
|
|
original = (
|
|
"[profiles.default]\n"
|
|
'model = "gpt-4o"\n'
|
|
"\n"
|
|
"[model_providers.openai]\n"
|
|
'base_url = "https://api.openai.com/v1"\n'
|
|
)
|
|
config_file.write_text(original)
|
|
|
|
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
|
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
|
assert wrap_result.exit_code == 0, wrap_result.output
|
|
assert 'model_provider = "headroom"' in config_file.read_text()
|
|
|
|
unwrap_result = runner.invoke(main, ["unwrap", "codex"])
|
|
assert unwrap_result.exit_code == 0, unwrap_result.output
|
|
|
|
# Config must be byte-for-byte what the user had before wrap, and the
|
|
# injected block must be gone — no more "Missing OPENAI_API_KEY" when the
|
|
# proxy is stopped.
|
|
assert config_file.read_text() == original
|
|
assert 'model_provider = "headroom"' not in config_file.read_text()
|
|
assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists()
|
|
|
|
|
|
def test_unwrap_codex_is_safe_noop_with_no_prior_wrap(
|
|
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
result = runner.invoke(main, ["unwrap", "codex"])
|
|
assert result.exit_code == 0, result.output
|
|
assert "Nothing to undo" in result.output
|
|
assert not (tmp_path / ".codex" / "config.toml").exists()
|
|
|
|
|
|
def test_unwrap_codex_removes_headroom_only_config_file(
|
|
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
|
|
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
|
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
|
assert wrap_result.exit_code == 0, wrap_result.output
|
|
|
|
config_file = tmp_path / ".codex" / "config.toml"
|
|
assert config_file.exists()
|
|
|
|
unwrap_result = runner.invoke(main, ["unwrap", "codex"])
|
|
assert unwrap_result.exit_code == 0, unwrap_result.output
|
|
assert not config_file.exists()
|
|
|
|
|
|
def test_unwrap_codex_preserves_unrelated_sections(
|
|
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
_set_test_home(monkeypatch, tmp_path)
|
|
config_file = tmp_path / ".codex" / "config.toml"
|
|
config_file.parent.mkdir(parents=True)
|
|
# A config with an MCP server the user configured by hand.
|
|
original = '[mcp_servers.local_thing]\ncommand = "/usr/local/bin/thing"\nargs = ["--serve"]\n'
|
|
config_file.write_text(original)
|
|
|
|
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
|
runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
|
|
|
result = runner.invoke(main, ["unwrap", "codex"])
|
|
assert result.exit_code == 0, result.output
|
|
restored = config_file.read_text()
|
|
assert restored == original
|