Spaces:
Running
Running
| """Regression tests for scripts/grpo_pool.py. | |
| Focus: `GrpoPool.connect()` must be all-or-nothing. If any single WebSocket | |
| handshake fails, every session that DID connect must be closed before the | |
| error propagates, so the server never ends up with leaked pool slots. | |
| No pytest-asyncio dependency — each test drives the loop via asyncio.run(). | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import pytest | |
| from scripts.grpo_pool import GrpoPool | |
| class _FakeEnv: | |
| """Minimal stand-in for AwsRlEnv. Tracks connect/close lifecycle.""" | |
| connect_calls = 0 # class-level so the factory can index envs in order | |
| def __init__(self, *, should_fail_on_index: int | None = None) -> None: | |
| self.connected = False | |
| self.close_called = False | |
| self._index = _FakeEnv.connect_calls | |
| _FakeEnv.connect_calls += 1 | |
| self._should_fail = ( | |
| should_fail_on_index is not None and self._index == should_fail_on_index | |
| ) | |
| async def connect(self) -> None: | |
| if self._should_fail: | |
| raise ConnectionError(f"fake failure on env#{self._index}") | |
| await asyncio.sleep(0) # yield so sibling connects can interleave | |
| self.connected = True | |
| async def close(self) -> None: | |
| self.close_called = True | |
| def _install_fake_env(monkeypatch, fail_on_index: int | None) -> list[_FakeEnv]: | |
| """Monkeypatch AwsRlEnv inside scripts.grpo_pool so GrpoPool builds FakeEnvs. | |
| Returns a shared list the test can inspect after connect() runs. | |
| """ | |
| _FakeEnv.connect_calls = 0 | |
| created: list[_FakeEnv] = [] | |
| def factory(*args, **kwargs) -> _FakeEnv: | |
| env = _FakeEnv(should_fail_on_index=fail_on_index) | |
| created.append(env) | |
| return env | |
| monkeypatch.setattr("scripts.grpo_pool.AwsRlEnv", factory) | |
| return created | |
| # --------------------------------------------------------------------------- | |
| # Happy path — sanity check the fake harness before running the failure cases | |
| # --------------------------------------------------------------------------- | |
| class TestConnectHappyPath: | |
| def test_all_sessions_connect_and_land_on_pool(self, monkeypatch) -> None: | |
| created = _install_fake_env(monkeypatch, fail_on_index=None) | |
| pool = GrpoPool(base_url="http://x", size=4) | |
| asyncio.run(pool.connect()) | |
| assert len(pool.envs) == 4 | |
| assert all(e.connected for e in created) | |
| assert not any(e.close_called for e in created) | |
| # --------------------------------------------------------------------------- | |
| # The review: partial failure must roll back | |
| # --------------------------------------------------------------------------- | |
| class TestConnectRollbackOnPartialFailure: | |
| def test_failure_closes_every_env_including_successful_ones( | |
| self, monkeypatch | |
| ) -> None: | |
| created = _install_fake_env(monkeypatch, fail_on_index=2) | |
| pool = GrpoPool(base_url="http://x", size=4) | |
| with pytest.raises(ConnectionError): | |
| asyncio.run(pool.connect()) | |
| # Every FakeEnv must have had close() called — successful ones so | |
| # server slots are released; the failing one as a harmless no-op. | |
| assert all(e.close_called for e in created), ( | |
| "Regression: successful sessions leaked after partial connect failure" | |
| ) | |
| def test_pool_envs_stays_empty_on_failure(self, monkeypatch) -> None: | |
| _install_fake_env(monkeypatch, fail_on_index=1) | |
| pool = GrpoPool(base_url="http://x", size=3) | |
| with pytest.raises(ConnectionError): | |
| asyncio.run(pool.connect()) | |
| # connect() must NOT leave a half-initialised pool visible to callers. | |
| assert pool.envs == [] | |
| def test_failure_does_not_block_retry(self, monkeypatch) -> None: | |
| """After a failed connect(), the caller can fix the root cause and | |
| call connect() again. pool.envs should be fresh.""" | |
| _install_fake_env(monkeypatch, fail_on_index=0) | |
| pool = GrpoPool(base_url="http://x", size=2) | |
| with pytest.raises(ConnectionError): | |
| asyncio.run(pool.connect()) | |
| # Second attempt with no injected failure should succeed. | |
| _install_fake_env(monkeypatch, fail_on_index=None) | |
| asyncio.run(pool.connect()) | |
| assert len(pool.envs) == 2 | |
| assert all(e.connected for e in pool.envs) | |
| def test_async_context_manager_cleans_up_when_enter_fails( | |
| self, monkeypatch | |
| ) -> None: | |
| """If `async with GrpoPool(...)` raises during __aenter__, | |
| __aexit__ is NOT called — so rollback must live inside connect() | |
| itself. This test exercises exactly that scenario. | |
| """ | |
| created = _install_fake_env(monkeypatch, fail_on_index=2) | |
| async def enter_and_fail() -> None: | |
| async with GrpoPool(base_url="http://x", size=4): | |
| pytest.fail("should never enter the body") | |
| with pytest.raises(ConnectionError): | |
| asyncio.run(enter_and_fail()) | |
| assert all(e.close_called for e in created) | |