diff --git a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py index 65f40770b1..182c946258 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py @@ -4,6 +4,8 @@ import tempfile import threading import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager from datetime import UTC, datetime, timedelta from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path @@ -89,6 +91,25 @@ def _event(stage_id: str = "compile_targets") -> dict[str, object]: } +@contextmanager +def _serving_thread( + target: Callable[[], object], stop: Callable[[], object] +) -> Iterator[threading.Thread]: + """Run a test server on a daemon thread; stop and join it however the block exits. + + Interpreter exit waits for every non-daemon thread, so a server thread left + running by a failed assertion hangs pytest after it prints its summary. + """ + + thread = threading.Thread(target=target, daemon=True) + thread.start() + try: + yield thread + finally: + stop() + thread.join(timeout=30) + + def test_development_collector_must_be_on_loopback() -> None: assert collector_module._development_collector_url("http://127.0.0.1:8080") == ( "http://127.0.0.1:8080" @@ -129,19 +150,15 @@ def do_POST(self): def log_message(self, format, *args): return - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server_thread = threading.Thread(target=server.serve_forever) - server_thread.start() - try: + with ( + ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server, + _serving_thread(server.serve_forever, server.shutdown), + ): status, _ = collector_module._http_post( f"http://127.0.0.1:{server.server_port}/exchange", {"run_id": "run-a"}, "hf-private-token", ) - finally: - server.shutdown() - server_thread.join(timeout=2) - server.server_close() assert status == 307 assert paths == ["/exchange"] @@ -604,27 +621,26 @@ def test_local_socket_acknowledges_after_durable_queue( heartbeat_seconds=60, drain_seconds=0, ) - thread = threading.Thread(target=emitter.run) - thread.start() - deadline = time.monotonic() + 2 - while not socket_path.exists() and time.monotonic() < deadline: - time.sleep(0.01) - - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: - client.connect(str(socket_path)) - client.sendall( - json.dumps({"action": "event", "event": _event()}).encode() + b"\n" - ) - assert client.recv(16) == b"ok\n" + with _serving_thread(emitter.run, emitter._stop.set) as thread: + deadline = time.monotonic() + 2 + while not socket_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(socket_path)) + client.sendall( + json.dumps({"action": "event", "event": _event()}).encode() + b"\n" + ) + assert client.recv(16) == b"ok\n" - assert spool.batch("run-a", "producer-a")[0]["resources"]["rss_bytes"] == 100 + assert spool.batch("run-a", "producer-a")[0]["resources"]["rss_bytes"] == 100 - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: - client.connect(str(socket_path)) - client.sendall(b'{"action":"close"}\n') - assert client.recv(16) == b"ok\n" - thread.join(timeout=2) - assert not thread.is_alive() + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(socket_path)) + client.sendall(b'{"action":"close"}\n') + assert client.recv(16) == b"ok\n" + thread.join(timeout=2) + assert not thread.is_alive() def test_subprocess_exchanges_ambient_token_and_delivers_events( @@ -659,26 +675,31 @@ def do_POST(self): def log_message(self, format, *args): return - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server_thread = threading.Thread(target=server.serve_forever) - server_thread.start() monkeypatch.setenv("HF_TOKEN", "hf-ambient-test-token") - emitter = LocalTelemetryEmitter.start( - run_id="subprocess-run", - country_code="US", - pipeline="test-pipeline", - development_collector_url=f"http://127.0.0.1:{server.server_port}", - spool_path=tmp_path / "subprocess.sqlite3", - heartbeat_seconds=60, - ) - assert emitter.available - emitter.stage("compile", message="Started.") - emitter.complete() - assert emitter._process is not None - emitter._process.wait(timeout=10) - server.shutdown() - server_thread.join(timeout=2) - server.server_close() + with ( + ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server, + _serving_thread(server.serve_forever, server.shutdown), + ): + emitter = LocalTelemetryEmitter.start( + run_id="subprocess-run", + country_code="US", + pipeline="test-pipeline", + development_collector_url=f"http://127.0.0.1:{server.server_port}", + spool_path=tmp_path / "subprocess.sqlite3", + heartbeat_seconds=60, + ) + try: + assert emitter.available + emitter.stage("compile", message="Started.") + emitter.complete() + assert emitter._process is not None + emitter._process.wait(timeout=10) + finally: + # A failure above can leave the service running; reap it before + # the collector it posts to shuts down. + if emitter._process is not None and emitter._process.poll() is None: + emitter._process.kill() + emitter._process.wait(timeout=10) assert requests[0][0] == "/v1/auth/huggingface/exchange" assert requests[0][1] == "Bearer hf-ambient-test-token" @@ -702,6 +723,45 @@ def log_message(self, format, *args): assert events[-1]["status"] == "completed" +def test_http_serving_thread_stops_when_the_test_body_fails() -> None: + """A failed assertion must not leave a server thread that blocks pytest's exit.""" + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format, *args): + return + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + with pytest.raises(AssertionError, match="forced failure"): + with _serving_thread(server.serve_forever, server.shutdown) as thread: + raise AssertionError("forced failure") + + assert thread.daemon + assert not thread.is_alive() + + +def test_emitter_serving_thread_stops_when_the_test_body_fails(tmp_path) -> None: + socket_path = ( + Path(tempfile.mkdtemp(prefix="microcosm-test-", dir="/tmp")) / "e.sock" + ) + emitter = EmitterService( + socket_path=socket_path, + registration=_registration(), + spool=EventSpool(tmp_path / "events.sqlite3"), + delivery=_FakeDelivery(), + sampler=_FakeSampler(), + heartbeat_seconds=60, + drain_seconds=0, + ) + + with pytest.raises(AssertionError, match="forced failure"): + with _serving_thread(emitter.run, emitter._stop.set) as thread: + raise AssertionError("forced failure") + + assert thread.daemon + assert not thread.is_alive() + assert not socket_path.exists() + + def test_failure_marks_active_stage_failed_before_run_failure(): from test_support.microcosm_build.telemetry import FakeTelemetryEmitter