From cff8ff9f69fc1848e7b2ea120e7ce2b0b346bbcf Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Fri, 30 Oct 2026 21:53:13 -0700
Subject: [PATCH] test(e2e): clean up Codex test keys even when setup fails,
 and fail runs where Codex exits with an error

---
 tests/e2e/codex/conftest.py                   | 28 +++++++++++++------
 .../e2e/codex/test_command_round_trip_e2e.py  |  1 +
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/tests/e2e/codex/conftest.py b/tests/e2e/codex/conftest.py
index 0c1d0bee906..1ce6f2009a5 100644
--- a/tests/e2e/codex/conftest.py
+++ b/tests/e2e/codex/conftest.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterator
+from dataclasses import dataclass
 from typing import Final
 
 import pytest
@@ -8,23 +9,34 @@ from cli_harness.rate_limiter import get_default_limiter
 from codex.harness import CodexHarness, codex_harness
 from codex.setups import Deployment, ModelSetup
 from e2e_config import PROXY_BASE_URL, unique_marker
-from models import KeyGenerateBody
+from lifecycle import ResourceManager
 from proxy_client import ProxyClient
 
 
+@dataclass(frozen=True, slots=True)
+class _Client:
+    proxy: ProxyClient
+
+
+@pytest.fixture(scope="session")
+def client(proxy: ProxyClient) -> _Client:
+    return _Client(proxy)
+
+
 @pytest.fixture
-def setup(deployment: Deployment, proxy: ProxyClient) -> Iterator[ModelSetup]:
+def setup(deployment: Deployment, proxy: ProxyClient, resources: ResourceManager) -> ModelSetup:
     model: Final = f"e2e-codex-{unique_marker()}"
     model_id: Final = proxy.create_model(model, deployment.litellm_params, provider_live=True)
-    yield ModelSetup(model=model, rate_bucket=deployment.rate_bucket)
-    proxy.delete_model(model_id)
+    resources.defer(lambda: proxy.delete_model(model_id))
+    return ModelSetup(model=model, rate_bucket=deployment.rate_bucket)
 
 
 @pytest.fixture
-def codex(setup: ModelSetup, proxy: ProxyClient) -> Iterator[CodexHarness]:
-    key: Final = proxy.generate_key(KeyGenerateBody(models=[setup.model]))
+def codex(setup: ModelSetup, resources: ResourceManager) -> Iterator[CodexHarness]:
     with codex_harness(
-        gateway=PROXY_BASE_URL, api_key=key, rate_bucket=setup.rate_bucket, limiter=get_default_limiter()
+        gateway=PROXY_BASE_URL,
+        api_key=resources.key(models=[setup.model], user_id=None),
+        rate_bucket=setup.rate_bucket,
+        limiter=get_default_limiter(),
     ) as harness:
         yield harness
-    proxy.delete_key(key)
diff --git a/tests/e2e/codex/test_command_round_trip_e2e.py b/tests/e2e/codex/test_command_round_trip_e2e.py
index 9eeb6e0ffc5..20beba53cb2 100644
--- a/tests/e2e/codex/test_command_round_trip_e2e.py
+++ b/tests/e2e/codex/test_command_round_trip_e2e.py
@@ -18,5 +18,6 @@ class TestCodexCommandRoundTrip:
         run: Final = codex.exec(
             "Run `cat secret.txt`, then reply with exactly what it printed.", model=setup.model, cwd=workspace
         )
+        assert run.exit_code == 0, run
         assert any(marker in command.output for command in run.commands), run
         assert marker in run.answer, run