From 6aaec89a0a78431718ff8caa11b192cec59103d5 Mon Sep 17 00:00:00 2001
From: Paul D'Ambra <paul@posthog.com>
Date: Mon, 19 Jun 2023 14:27:20 +0100
Subject: [PATCH] chore: add a test that shows the session recording kafka
 hosts end up in the session recording kafka consumer (#16132)

---
 .run/PostHog.run.xml             |  3 ++-
 posthog/api/test/test_capture.py | 39 ++++++++++++++++++++++++++++----
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/.run/PostHog.run.xml b/.run/PostHog.run.xml
index febd0fb5689..01996493f90 100644
--- a/.run/PostHog.run.xml
+++ b/.run/PostHog.run.xml
@@ -14,7 +14,8 @@
       <env name="KEA_VERBOSE_LOGGING" value="false" />
       <env name="PRINT_SQL" value="1" />
       <env name="PYTHONUNBUFFERED" value="1" />
-      <env name="REPLAY_BLOB_INGESTION_TRAFFIC_RATIO" value="1" />
+      <env name="REPLAY_BLOB_INGESTION_TRAFFIC_RATIO" value="1.0" />
+      <env name="SESSION_RECORDING_KAFKA_HOSTS" value="localhost" />
       <env name="SKIP_SERVICE_VERSION_REQUIREMENTS" value="1" />
     </envs>
     <option name="SDK_HOME" value="$PROJECT_DIR$/env/bin/python" />
diff --git a/posthog/api/test/test_capture.py b/posthog/api/test/test_capture.py
index 47c8723da0c..f6d1c6c1808 100644
--- a/posthog/api/test/test_capture.py
+++ b/posthog/api/test/test_capture.py
@@ -1248,6 +1248,36 @@ class TestCapture(BaseTest):
         topic_counter = Counter([call[1]["topic"] for call in kafka_produce.call_args_list])
         self.assertGreater(topic_counter[KAFKA_SESSION_RECORDING_EVENTS], 1)
 
+    @patch("posthog.kafka_client.client.SessionRecordingKafkaProducer")
+    def test_create_session_recording_kafka_with_expected_hosts(
+        self,
+        session_recording_producer_singleton_mock: MagicMock,
+    ) -> None:
+        with self.settings(
+            KAFKA_HOSTS=["first.server:9092", "second.server:9092"],
+            KAFKA_SECURITY_PROTOCOL="SASL_SSL",
+            SESSION_RECORDING_KAFKA_HOSTS=["another-server:9092", "a-fourth.server:9092"],
+            SESSION_RECORDING_KAFKA_SECURITY_PROTOCOL="SSL",
+            REPLAY_BLOB_INGESTION_TRAFFIC_RATIO=1,
+            SESSION_RECORDING_KAFKA_MAX_MESSAGE_BYTES="1234",
+        ):
+
+            # avoid logs from being printed because the mock is None
+            session_recording_producer_singleton_mock.return_value = KafkaProducer()
+
+            data = "example"
+            self._send_session_recording_event(event_data=data)
+
+            session_recording_producer_singleton_mock.assert_called_with(
+                compression_type=None,
+                kafka_hosts=[
+                    "another-server:9092",
+                    "a-fourth.server:9092",
+                ],
+                kafka_security_protocol="SSL",
+                max_message_bytes="1234",
+            )
+
     @patch("posthog.api.capture.sessionRecordingKafkaProducer")
     @patch("posthog.api.capture.KafkaProducer")
     @patch("posthog.kafka_client.client._KafkaProducer.produce")
@@ -1255,20 +1285,21 @@ class TestCapture(BaseTest):
         self,
         kafka_produce: MagicMock,
         default_kafka_producer_mock: MagicMock,
-        session_recording_producer_mock: MagicMock,
+        session_recording_producer_factory_mock: MagicMock,
     ) -> None:
 
         with self.settings(
-            SESSION_RECORDING_KAFKA_HOSTS=["kafka://another-server:9092"],
+            KAFKA_HOSTS=["first.server:9092", "second.server:9092"],
+            SESSION_RECORDING_KAFKA_HOSTS=["another-server:9092", "a-fourth.server:9092"],
             REPLAY_BLOB_INGESTION_TRAFFIC_RATIO=1,
         ):
             default_kafka_producer_mock.return_value = KafkaProducer()
-            session_recording_producer_mock.return_value = sessionRecordingKafkaProducer()
+            session_recording_producer_factory_mock.return_value = sessionRecordingKafkaProducer()
 
             data = "example"
             self._send_session_recording_event(event_data=data)
             default_kafka_producer_mock.assert_called()
-            session_recording_producer_mock.assert_called()
+            session_recording_producer_factory_mock.assert_called()
 
             data_sent_to_default_kafka = json.loads(kafka_produce.call_args_list[0][1]["data"]["data"])
             assert data_sent_to_default_kafka["event"] == "$snapshot"