test: capture invariant hardening

Author: eli-r-ph

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]

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

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

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)