From eff6dc71aa3851ea2a8d70c0cc6a0907dcda19d1 Mon Sep 17 00:00:00 2001
From: Xavier Vello <xavier.vello@gmail.com>
Date: Thu, 26 Oct 2023 15:31:48 +0200
Subject: [PATCH] feat(decide): add NEW_ANALYTICS_CAPTURE_SAMPLING_RATE option
 for slower rollout (#18203)

---
 posthog/api/decide.py           |  3 +-
 posthog/api/test/test_decide.py | 58 +++++++++++++++++++++++----------
 posthog/settings/ingestion.py   |  1 +
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/posthog/api/decide.py b/posthog/api/decide.py
index 1f173bdc39016..0991cacd7d9a7 100644
--- a/posthog/api/decide.py
+++ b/posthog/api/decide.py
@@ -222,7 +222,8 @@ def get_decide(request: HttpRequest):
             )
 
             if settings.NEW_ANALYTICS_CAPTURE_TEAM_IDS and str(team.id) in settings.NEW_ANALYTICS_CAPTURE_TEAM_IDS:
-                response["analytics"] = {"endpoint": settings.NEW_ANALYTICS_CAPTURE_ENDPOINT}
+                if random() < settings.NEW_ANALYTICS_CAPTURE_SAMPLING_RATE:
+                    response["analytics"] = {"endpoint": settings.NEW_ANALYTICS_CAPTURE_ENDPOINT}
 
             if team.session_recording_opt_in and (
                 on_permitted_recording_domain(team, request) or not team.recording_domains
diff --git a/posthog/api/test/test_decide.py b/posthog/api/test/test_decide.py
index 78019779fb5f1..804e8e8d86460 100644
--- a/posthog/api/test/test_decide.py
+++ b/posthog/api/test/test_decide.py
@@ -1,25 +1,23 @@
 import base64
 import json
 import random
-from unittest.mock import patch
 import time
-from django.conf import settings
-
+from unittest.mock import patch
 
+import pytest
+from django.conf import settings
 from django.core.cache import cache
 from django.db import connection, connections
 from django.test import TransactionTestCase, TestCase
 from django.test.client import Client
-from rest_framework.test import APIClient
 from freezegun import freeze_time
-import pytest
 from rest_framework import status
-from posthog.models.feature_flag.feature_flag import FeatureFlagHashKeyOverride
-from posthog.models.group.group import Group
-
+from rest_framework.test import APIClient
 
-from posthog.api.test.test_feature_flag import QueryTimeoutWrapper
+from posthog import redis
 from posthog.api.decide import label_for_team_id_to_track
+from posthog.api.test.test_feature_flag import QueryTimeoutWrapper
+from posthog.database_healthcheck import postgres_healthcheck
 from posthog.models import (
     FeatureFlag,
     GroupTypeMapping,
@@ -30,21 +28,16 @@
     PluginSourceFile,
 )
 from posthog.models.cohort.cohort import Cohort
+from posthog.models.feature_flag.feature_flag import FeatureFlagHashKeyOverride
+from posthog.models.group.group import Group
 from posthog.models.organization import Organization, OrganizationMembership
+from posthog.models.person import PersonDistinctId
 from posthog.models.personal_api_key import hash_key_value
 from posthog.models.plugin import sync_team_inject_web_apps
 from posthog.models.team.team import Team
-from posthog.models.person import PersonDistinctId
 from posthog.models.user import User
 from posthog.models.utils import generate_random_token_personal
-from posthog.test.base import (
-    BaseTest,
-    QueryMatchingTest,
-    snapshot_postgres_queries,
-    snapshot_postgres_queries_context,
-)
-from posthog.database_healthcheck import postgres_healthcheck
-from posthog import redis
+from posthog.test.base import BaseTest, QueryMatchingTest, snapshot_postgres_queries, snapshot_postgres_queries_context
 
 
 @patch(
@@ -2943,6 +2936,35 @@ def test_decide_analytics_fire_for_survey_targeting_flags(self, *args):
                 {b"165192618": b"1"},
             )
 
+    @patch("posthog.models.feature_flag.flag_analytics.CACHE_BUCKET_SIZE", 10)
+    def test_decide_new_capture_activation(self, *args):
+        self.client.logout()
+        with self.settings(NEW_ANALYTICS_CAPTURE_TEAM_IDS={str(self.team.id)}, NEW_ANALYTICS_CAPTURE_SAMPLING_RATE=1.0):
+            response = self._post_decide(api_version=3)
+            self.assertEqual(response.status_code, 200)
+            self.assertTrue("analytics" in response.json())
+            self.assertEqual(response.json()["analytics"]["endpoint"], "/i/v0/e/")
+
+        with self.settings(
+            NEW_ANALYTICS_CAPTURE_TEAM_IDS={str(self.team.id)},
+            NEW_ANALYTICS_CAPTURE_SAMPLING_RATE=1.0,
+            NEW_ANALYTICS_CAPTURE_ENDPOINT="/custom",
+        ):
+            response = self._post_decide(api_version=3)
+            self.assertEqual(response.status_code, 200)
+            self.assertTrue("analytics" in response.json())
+            self.assertEqual(response.json()["analytics"]["endpoint"], "/custom")
+
+        with self.settings(NEW_ANALYTICS_CAPTURE_TEAM_IDS={"0"}, NEW_ANALYTICS_CAPTURE_SAMPLING_RATE=1.0):
+            response = self._post_decide(api_version=3)
+            self.assertEqual(response.status_code, 200)
+            self.assertFalse("analytics" in response.json())
+
+        with self.settings(NEW_ANALYTICS_CAPTURE_TEAM_IDS={str(self.team.id)}, NEW_ANALYTICS_CAPTURE_SAMPLING_RATE=0):
+            response = self._post_decide(api_version=3)
+            self.assertEqual(response.status_code, 200)
+            self.assertFalse("analytics" in response.json())
+
 
 class TestDatabaseCheckForDecide(BaseTest, QueryMatchingTest):
     """
diff --git a/posthog/settings/ingestion.py b/posthog/settings/ingestion.py
index b60206d101ae9..bab3d7f2506ae 100644
--- a/posthog/settings/ingestion.py
+++ b/posthog/settings/ingestion.py
@@ -36,3 +36,4 @@
 
 NEW_ANALYTICS_CAPTURE_ENDPOINT = os.getenv("NEW_CAPTURE_ENDPOINT", "/i/v0/e/")
 NEW_ANALYTICS_CAPTURE_TEAM_IDS = get_set(os.getenv("NEW_ANALYTICS_CAPTURE_TEAM_IDS", ""))
+NEW_ANALYTICS_CAPTURE_SAMPLING_RATE = get_from_env("NEW_ANALYTICS_CAPTURE_SAMPLING_RATE", type_cast=float, default=1.0)