From d28b5beb200353f0ca6ef561513091307b75b9fc Mon Sep 17 00:00:00 2001
From: Ted Kaemming <65315+tkaemming@users.noreply.github.com>
Date: Tue, 26 Mar 2024 16:24:06 -0700
Subject: [PATCH] Extend capture test to cover setting enabled and disabled

---
 posthog/api/test/test_capture.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/posthog/api/test/test_capture.py b/posthog/api/test/test_capture.py
index 2a80186082dea..a0fc8826c95c6 100644
--- a/posthog/api/test/test_capture.py
+++ b/posthog/api/test/test_capture.py
@@ -4,6 +4,7 @@
 import base64
 import gzip
 import json
+from django.test import override_settings
 import lzstring
 import pathlib
 import pytest
@@ -281,8 +282,7 @@ def test_is_randomly_partitioned(self):
             assert is_randomly_partitioned(override_key) is True
 
     @patch("posthog.kafka_client.client._KafkaProducer.produce")
-    def test_capture_randomly_partitions_with_likely_anonymous_ids(self, kafka_produce):
-        """Test is_randomly_partitioned in the prescence of likely anonymous ids."""
+    def _do_test_capture_with_likely_anonymous_ids(self, kafka_produce, expect_random_partitioning: bool):
         for distinct_id in LIKELY_ANONYMOUS_IDS:
             data = {
                 "event": "$autocapture",
@@ -298,9 +298,23 @@ def test_capture_randomly_partitions_with_likely_anonymous_ids(self, kafka_produ
                 )
 
             kafka_produce.assert_called_with(
-                topic=KAFKA_EVENTS_PLUGIN_INGESTION_TOPIC, data=ANY, key=None, headers=None
+                topic=KAFKA_EVENTS_PLUGIN_INGESTION_TOPIC,
+                data=ANY,
+                key=None if expect_random_partitioning else ANY,
+                headers=None,
             )
 
+            if not expect_random_partitioning:
+                assert kafka_produce.mock_calls[0].kwargs["key"] is not None
+
+    def test_capture_randomly_partitions_with_likely_anonymous_ids(self):
+        """Test is_randomly_partitioned in the prescence of likely anonymous ids, if enabled."""
+        with override_settings(CAPTURE_ALLOW_RANDOM_PARTITIONING=True):
+            self._do_test_capture_with_likely_anonymous_ids(expect_random_partitioning=True)
+
+        with override_settings(CAPTURE_ALLOW_RANDOM_PARTITIONING=False):
+            self._do_test_capture_with_likely_anonymous_ids(expect_random_partitioning=False)
+
     def test_cached_is_randomly_partitioned(self):
         """Assert the behavior of is_randomly_partitioned under certain cache settings.