From 2011a51b0206340c5a77a584eb15458cfef77563 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 6 Oct 2026 20:43:27 -0700 Subject: [PATCH 1/2] test: capture invariant hardening --- posthog/test/test_ai_capture_lane.py | 12 +++ posthog/test/test_async_client.py | 47 ++++++---- posthog/test/test_capture_invariants.py | 111 ++++++++++++++++++++++++ posthog/test/test_client.py | 50 +++++++++-- 4 files changed, 195 insertions(+), 25 deletions(-) create mode 100644 posthog/test/test_capture_invariants.py diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 885d4440..d3ff8762 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -209,6 +209,18 @@ def test_sync_mode_ai_event_over_cap_is_not_sent(self): mock_send.assert_not_called() self.assertIn("exceeds the 1024KiB limit", "\n".join(logs.output)) + def test_sync_mode_ai_event_at_endpoint_ceiling_is_sent(self): + client = Client(TEST_API_KEY, sync_mode=True) + with patch_capture_send("client") as mock_send: + result = client.capture_ai( + "$ai_generation", + distinct_id="d", + properties={"p": "x" * AI_MAX_PROPERTIES_SIZE}, + ) + + self.assertIsNotNone(result) + mock_send.assert_called_once() + def test_analytics_lane_rejects_events_over_900kib(self): client = self._client() consumer = client.consumers[0] diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 2a5af68e..d6f84950 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -13,7 +13,7 @@ from posthog import AsyncClient, AsyncPosthog, CaptureCompression from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH -from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE +from posthog.consumer import AI_MAX_MSG_SIZE, AI_MAX_PROPERTIES_SIZE, MAX_MSG_SIZE from posthog.contexts import ( new_context, set_capture_exception_code_variables_context, @@ -123,7 +123,10 @@ async def send_batch(api_key, host, batch, **kwargs): @pytest.mark.asyncio -async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): +@pytest.mark.parametrize("method_name", ["capture", "capture_ai"]) +async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission( + method_name, +): client = AsyncPosthog("test-key", flush_interval=30) client._ensure_workers_started(client._analytics_lane) scheduled_callbacks = [] @@ -136,7 +139,7 @@ async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): ): capture_thread = threading.Thread( target=lambda: capture_result.append( - client.capture("threaded event", distinct_id="user-1") + getattr(client, method_name)("threaded event", distinct_id="user-1") ), daemon=True, ) @@ -149,9 +152,12 @@ async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): capture_thread.join(timeout=1) assert capture_result == [None] - for task in client._analytics_lane.worker_tasks: + assert client._ai_lane.worker_tasks == [] + assert client._pending_queue_items() == 0 + worker_tasks = client._all_worker_tasks() + for task in worker_tasks: task.cancel() - await asyncio.gather(*client._analytics_lane.worker_tasks, return_exceptions=True) + await asyncio.gather(*worker_tasks, return_exceptions=True) await client._close_transport() @@ -288,7 +294,8 @@ async def send_batch(api_key, host, batch, **kwargs): ("method_name", "client_kwargs", "payload_size", "sent"), [ ("capture_immediate", {}, MAX_MSG_SIZE, False), - ("capture_ai_immediate", {}, MAX_MSG_SIZE, True), + ("capture_ai_immediate", {}, AI_MAX_PROPERTIES_SIZE, True), + ("capture_ai_immediate", {}, AI_MAX_MSG_SIZE, False), ("capture_ai_immediate", {"capture_ai_max_event_bytes": 1024}, 2048, False), ], ) @@ -359,7 +366,7 @@ async def send_batch(api_key, host, batch, **kwargs): client.capture_ai( "$ai_generation", distinct_id="user-1", - properties={"$ai_input": "x" * MAX_MSG_SIZE}, + properties={"$ai_input": "x" * AI_MAX_PROPERTIES_SIZE}, ) await client.flush(timeout_seconds=1) assert sent == { @@ -800,23 +807,31 @@ async def on_error(error, batch): @pytest.mark.asyncio -@pytest.mark.parametrize("immediate", [False, True]) -async def test_failed_capture_does_not_log_server_response_detail(caplog, immediate): +@pytest.mark.parametrize( + ("method_name", "path"), + [ + ("capture", _CAPTURE_V1_PATH), + ("capture_immediate", _CAPTURE_V1_PATH), + ("capture_ai", _CAPTURE_AI_V1_PATH), + ("capture_ai_immediate", _CAPTURE_AI_V1_PATH), + ], +) +async def test_failed_capture_logs_one_loss_line_without_response_detail( + caplog, method_name, path +): caplog.set_level(logging.DEBUG, logger="posthog") server_error = APIError(400, "password=server-secret") with patch_async_capture_send(side_effect=server_error): client = AsyncPosthog("test-key", flush_at=1, max_retries=0) - if immediate: - await client.capture_immediate("event", distinct_id="user-1") - else: - client.capture("event", distinct_id="user-1") - await client.flush(timeout_seconds=1) + result = getattr(client, method_name)("$ai_generation", distinct_id="user-1") + if method_name.endswith("_immediate"): + await result + await client.flush(timeout_seconds=1) await client.shutdown() assert "server-secret" not in caplog.text - assert "APIError" in caplog.text - assert "status=400" in caplog.text + assert f"1 event(s) not persisted by {path}: APIError (status=400)" in caplog.text @pytest.mark.asyncio diff --git a/posthog/test/test_capture_invariants.py b/posthog/test/test_capture_invariants.py new file mode 100644 index 00000000..48f933be --- /dev/null +++ b/posthog/test/test_capture_invariants.py @@ -0,0 +1,111 @@ +import json +from typing import Any, Optional +from unittest import mock + +import pytest +from requests import Response + +from posthog import AsyncPosthog, CaptureCompression, Client +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH + +_UUID = "01890000-0000-7000-8000-000000000001" +_TIMESTAMP = "2026-01-02T03:04:05+00:00" + + +def _event_kwargs() -> dict[str, Any]: + return { + "distinct_id": "user-1", + "uuid": _UUID, + "timestamp": _TIMESTAMP, + "properties": { + "$process_person_profile": False, + "$product_tour_id": "tour-1", + "empty_list": [], + "empty_map": {}, + "none": None, + "nested": {"a": {"b": [1, None, {}]}}, + }, + "options": { + "process_person_profile": None, + "cookieless_mode": False, + "custom_list": [], + "custom_map": {"inner": {}}, + }, + } + + +def _record_wire(sent: list[tuple[str, list[dict]]]): + def post(url, data=None, **kwargs): + events = json.loads(data)["batch"] + sent.append((url, events)) + response = Response() + response.status_code = 200 + response._content = json.dumps( + {"results": {event["uuid"]: {"result": "ok"} for event in events}} + ).encode() + return response + + session = mock.Mock() + session.post.side_effect = post + return mock.patch("posthog.capture_send._get_session", return_value=session) + + +def _send_sync(method_name: str, before_send) -> list[tuple[str, list[dict]]]: + sent: list[tuple[str, list[dict]]] = [] + with _record_wire(sent): + client = Client( + "test-key", + sync_mode=True, + before_send=before_send, + capture_compression=CaptureCompression.NONE, + ) + assert getattr(client, method_name)("$ai_generation", **_event_kwargs()) + client.shutdown() + return sent + + +async def _send_async(method_name: str, before_send) -> list[tuple[str, list[dict]]]: + sent: list[tuple[str, list[dict]]] = [] + with _record_wire(sent): + client = AsyncPosthog( + "test-key", + before_send=before_send, + capture_compression=CaptureCompression.NONE, + ) + assert await getattr(client, method_name)("$ai_generation", **_event_kwargs()) + await client.shutdown() + return sent + + +def _pass_through(event: dict) -> Optional[dict]: + return event + + +@pytest.mark.parametrize( + ("method_name", "path"), + [("capture", _CAPTURE_V1_PATH), ("capture_ai", _CAPTURE_AI_V1_PATH)], +) +def test_pass_through_before_send_leaves_sync_wire_unchanged(method_name, path): + without_hook = _send_sync(method_name, None) + with_hook = _send_sync(method_name, _pass_through) + + assert [url.endswith(path) for url, _ in without_hook] == [True] + assert with_hook == without_hook + assert without_hook[0][1][0]["options"]["process_person_profile"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "path"), + [ + ("capture_immediate", _CAPTURE_V1_PATH), + ("capture_ai_immediate", _CAPTURE_AI_V1_PATH), + ], +) +async def test_pass_through_before_send_leaves_async_wire_unchanged(method_name, path): + without_hook = await _send_async(method_name, None) + with_hook = await _send_async(method_name, _pass_through) + + assert [url.endswith(path) for url, _ in without_hook] == [True] + assert with_hook == without_hook + assert without_hook[0][1][0]["options"]["process_person_profile"] is False diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index c2b956eb..c58afe61 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -18,7 +18,7 @@ import pytest from posthog.capture_compression import CaptureCompression -from posthog.capture_send import _CAPTURE_V1_PATH +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.client import Client from posthog.contexts import get_context_session_id, new_context, set_context_session from posthog.request import APIError, GetResponse @@ -2764,19 +2764,29 @@ def on_error(err, batch): "\n".join(logs.output), ) - def test_sync_send_failure_without_on_error_logs_aggregate_line(self): + @parameterized.expand( + [ + ("analytics", "capture", _CAPTURE_V1_PATH), + ("ai", "capture_ai", _CAPTURE_AI_V1_PATH), + ] + ) + def test_sync_send_failure_without_on_error_logs_aggregate_line( + self, _name, method_name, path + ): client = Client(FAKE_TEST_API_KEY, sync_mode=True) with patch_capture_send( "client", side_effect=APIError(400, "password=server-secret") ): with self.assertLogs("posthog", level="ERROR") as logs: - result = client.capture("event", distinct_id="distinct_id") + result = getattr(client, method_name)( + "$ai_generation", distinct_id="distinct_id" + ) self.assertIsNone(result) output = "\n".join(logs.output) self.assertIn( - "1 event(s) not persisted by /i/v1/analytics/events: APIError (status=400)", + f"1 event(s) not persisted by {path}: APIError (status=400)", output, ) self.assertNotIn("server-secret", output) @@ -3401,12 +3411,21 @@ def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): self.assertLess(time.monotonic() - start, 1) - def test_shutdown_waits_for_racing_enqueue_before_draining(self): + @parameterized.expand( + [ + ("analytics", "capture", "_analytics_lane", _CAPTURE_V1_PATH), + ("ai_first_event", "capture_ai", "_ai_lane", _CAPTURE_AI_V1_PATH), + ] + ) + def test_shutdown_waits_for_racing_enqueue_before_draining( + self, _name, method_name, lane_attr, path + ): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + lane = getattr(client, lane_attr) put_started = threading.Event() release_put = threading.Event() shutdown_done = threading.Event() - original_put = client.queue.put + original_put = lane.queue.put capture_result = [] def blocking_put(*args, **kwargs): @@ -3416,14 +3435,17 @@ def blocking_put(*args, **kwargs): capture_thread = threading.Thread( target=lambda: capture_result.append( - client.capture("racing event", distinct_id="distinct_id") + getattr(client, method_name)("$ai_racing", distinct_id="distinct_id") ) ) shutdown_thread = threading.Thread( target=lambda: (client.shutdown(), shutdown_done.set()) ) - with mock.patch.object(client.queue, "put", side_effect=blocking_put): + with ( + patch_capture_send("consumer") as mock_send, + mock.patch.object(lane.queue, "put", side_effect=blocking_put), + ): capture_thread.start() self.assertTrue(put_started.wait(2)) shutdown_thread.start() @@ -3439,7 +3461,17 @@ def blocking_put(*args, **kwargs): self.assertFalse(shutdown_thread.is_alive()) self.assertTrue(shutdown_done.is_set()) self.assertIsNotNone(capture_result[0]) - self.assertTrue(client.queue.empty()) + self.assertTrue(lane.queue.empty()) + self.assertEqual( + [ + (call.kwargs["path"], len(sent_batch(mock_send, i))) + for i, call in enumerate(mock_send.call_args_list) + ], + [(path, 1)], + ) + self.assertFalse( + any(c.is_alive() for lane in client._lanes for c in lane.consumers) + ) def test_shutdown_waits_for_sync_send_and_rejects_later_sends(self): client = Client(FAKE_TEST_API_KEY, sync_mode=True) From aac873519ed9d1966e9f7ef81b5b04cefb167db4 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Wed, 7 Oct 2026 17:06:42 -0700 Subject: [PATCH 2/2] test: assert exactly one capture loss line per failed batch --- posthog/test/test_async_client.py | 9 ++++++++- posthog/test/test_client.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index d6f84950..a9ae4d72 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -831,7 +831,14 @@ async def test_failed_capture_logs_one_loss_line_without_response_detail( await client.shutdown() assert "server-secret" not in caplog.text - assert f"1 event(s) not persisted by {path}: APIError (status=400)" in caplog.text + loss_lines = [ + r.getMessage() + for r in caplog.records + if "event(s) not persisted by" in r.getMessage() + ] + assert loss_lines == [ + f"[PostHog] 1 event(s) not persisted by {path}: APIError (status=400)" + ] @pytest.mark.asyncio diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index c58afe61..d846ab6c 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -2785,9 +2785,14 @@ def test_sync_send_failure_without_on_error_logs_aggregate_line( self.assertIsNone(result) output = "\n".join(logs.output) - self.assertIn( - f"1 event(s) not persisted by {path}: APIError (status=400)", - output, + loss_lines = [ + r.getMessage() + for r in logs.records + if "event(s) not persisted by" in r.getMessage() + ] + self.assertEqual( + loss_lines, + [f"[PostHog] 1 event(s) not persisted by {path}: APIError (status=400)"], ) self.assertNotIn("server-secret", output)