Have an option to send feature variants with the .capture(...) calls (#65)

Co-authored-by: Utku Zihnioglu <[email protected]>

Author: neilkakkar

example.py

@@ -11,10 +11,10 @@
 
 # Where you host PostHog, with no trailing /.
 # You can remove this line if you're using posthog.com
-posthog.host = "http://127.0.0.1:8000"
+posthog.host = "http://localhost:8000"
 
 # Capture an event
-posthog.capture("distinct_id", "event", {"property1": "value", "property2": "value"})
+posthog.capture("distinct_id", "event", {"property1": "value", "property2": "value"}, send_feature_flags=True)
 
 print(posthog.feature_enabled("beta-feature", "distinct_id"))
 print(posthog.feature_enabled("beta-feature", "distinct_id", groups={"company": "id:5"}))

posthog/__init__.py

@@ -27,6 +27,7 @@ def capture(
     timestamp=None,  # type: Optional[datetime.datetime]
     uuid=None,  # type: Optional[str]
     groups=None,  # type: Optional[Dict]
+    send_feature_flags=False,
 ):
     # type: (...) -> None
     """
@@ -58,6 +59,7 @@ def capture(
         timestamp=timestamp,
         uuid=uuid,
         groups=groups,
+        send_feature_flags=send_feature_flags,
     )
 
 

posthog/client.py

@@ -120,8 +120,33 @@ def identify(self, distinct_id=None, properties=None, context=None, timestamp=No
 
         return self._enqueue(msg)
 
+    def get_feature_variants(self, distinct_id, groups=None):
+        assert self.personal_api_key, "You have to specify a personal_api_key to use feature flags."
+        require("distinct_id", distinct_id, ID_TYPES)
+
+        if groups:
+            require("groups", groups, dict)
+        else:
+            groups = {}
+
+        request_data = {
+            "distinct_id": distinct_id,
+            "personal_api_key": self.personal_api_key,
+            "groups": groups,
+        }
+        resp_data = decide(self.api_key, self.host, timeout=10, **request_data)
+        return resp_data["featureFlags"]
+
     def capture(
-        self, distinct_id=None, event=None, properties=None, context=None, timestamp=None, uuid=None, groups=None
+        self,
+        distinct_id=None,
+        event=None,
+        properties=None,
+        context=None,
+        timestamp=None,
+        uuid=None,
+        groups=None,
+        send_feature_flags=False,
     ):
         properties = properties or {}
         context = context or {}
@@ -142,6 +167,16 @@ def capture(
             require("groups", groups, dict)
             msg["properties"]["$groups"] = groups
 
+        if send_feature_flags:
+            try:
+                feature_variants = self.get_feature_variants(distinct_id, groups)
+            except Exception as e:
+                self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
+            else:
+                for feature, variant in feature_variants.items():
+                    msg["properties"]["$feature/{}".format(feature)] = variant
+                msg["properties"]["$active_feature_flags"] = list(feature_variants.keys())
+
         return self._enqueue(msg)
 
     def set(self, distinct_id=None, properties=None, context=None, timestamp=None, uuid=None):
@@ -364,23 +399,20 @@ def feature_enabled(self, key, distinct_id, default=False, *, groups={}):
                 if flag["key"] == key:
                     feature_flag = flag
                     if feature_flag.get("is_simple_flag"):
-                        response = _hash(key, distinct_id) <= (feature_flag.get("rollout_percentage", 100) / 100)
+                        rollout_percentage = (
+                            feature_flag.get("rollout_percentage")
+                            if feature_flag.get("rollout_percentage") is not None
+                            else 100
+                        )
+                        response = _hash(key, distinct_id) <= (rollout_percentage / 100)
         if response == None:
             try:
-                request_data = {
-                    "distinct_id": distinct_id,
-                    "personal_api_key": self.personal_api_key,
-                    "groups": groups,
-                }
-                resp_data = decide(self.api_key, self.host, timeout=10, **request_data)
+                feature_flags = self.get_feature_variants(distinct_id, groups=groups)
             except Exception as e:
+                self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
                 response = default
-                self.log.warning(
-                    "[FEATURE FLAGS] Unable to get data for flag %s, because of the following error:" % key
-                )
-                self.log.warning(e)
             else:
-                response = resp_data["featureFlags"].get(key, default)
+                response = feature_flags.get(key, default)
         self.capture(distinct_id, "$feature_flag_called", {"$feature_flag": key, "$feature_flag_response": response})
         return response
 

posthog/test/test_client.py

@@ -1,6 +1,7 @@
 import time
 import unittest
 from datetime import date, datetime
+from unittest.mock import MagicMock
 from uuid import uuid4
 
 import mock
@@ -74,6 +75,71 @@ def test_basic_capture_with_project_api_key(self):
         self.assertEqual(msg["properties"]["$lib"], "posthog-python")
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
 
+    @mock.patch("posthog.client.decide")
+    def test_basic_capture_with_feature_flags(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+
+        client = Client(TEST_API_KEY, on_error=self.set_fail, personal_api_key=TEST_API_KEY)
+        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True)
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
+        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature"])
+
+        self.assertEqual(patch_decide.call_count, 1)
+
+    @mock.patch("posthog.client.decide")
+    def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+
+        client = Client(TEST_API_KEY, on_error=self.set_fail, personal_api_key=TEST_API_KEY)
+        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=False)
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertTrue("$feature/beta-feature" not in msg["properties"])
+        self.assertTrue("$active_feature_flags" not in msg["properties"])
+
+        self.assertEqual(patch_decide.call_count, 0)
+
+    @mock.patch("posthog.client.decide")
+    def test_basic_capture_with_feature_flags_without_api_key(self, patch_decide):
+        patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}
+
+        client = Client(project_api_key=TEST_API_KEY, on_error=self.set_fail)
+        client.log = MagicMock()
+        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True)
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+
+        self.assertEqual(client.log.exception.call_count, 1)
+        client.log.exception.assert_called_with(
+            "[FEATURE FLAGS] Unable to get feature variants: You have to specify a personal_api_key to use feature flags."
+        )
+
     def test_stringifies_distinct_id(self):
         # A large number that loses precision in node:
         # node -e "console.log(157963456373623802 + 1)" > 157963456373623800
@@ -424,6 +490,16 @@ def test_feature_enabled_simple_is_false(self, patch_get, patch_decide):
         self.assertFalse(client.feature_enabled("beta-feature", "distinct_id"))
         self.assertEqual(patch_decide.call_count, 0)
 
+    @mock.patch("posthog.client.decide")
+    @mock.patch("posthog.client.get")
+    def test_feature_enabled_simple_is_true_when_rollout_is_undefined(self, patch_get, patch_decide):
+        client = Client(TEST_API_KEY)
+        client.feature_flags = [
+            {"id": 1, "name": "Beta Feature", "key": "beta-feature", "is_simple_flag": True, "rollout_percentage": None}
+        ]
+        self.assertTrue(client.feature_enabled("beta-feature", "distinct_id"))
+        self.assertEqual(patch_decide.call_count, 0)
+
     @mock.patch("posthog.client.get")
     def test_feature_enabled_simple_with_project_api_key(self, patch_get):
         client = Client(project_api_key=TEST_API_KEY, on_error=self.set_fail)
@@ -435,7 +511,7 @@ def test_feature_enabled_simple_with_project_api_key(self, patch_get):
     @mock.patch("posthog.client.decide")
     def test_feature_enabled_request(self, patch_decide):
         patch_decide.return_value = {"featureFlags": {"beta-feature": True}}
-        client = Client(TEST_API_KEY)
+        client = Client(TEST_API_KEY, personal_api_key="test")
         client.feature_flags = [
             {"id": 1, "name": "Beta Feature", "key": "beta-feature", "is_simple_flag": False, "rollout_percentage": 100}
         ]
@@ -444,7 +520,7 @@ def test_feature_enabled_request(self, patch_decide):
     @mock.patch("posthog.client.decide")
     def test_feature_enabled_request_multi_variate(self, patch_decide):
         patch_decide.return_value = {"featureFlags": {"beta-feature": "variant-1"}}
-        client = Client(TEST_API_KEY)
+        client = Client(TEST_API_KEY, personal_api_key="test")
         client.feature_flags = [
             {"id": 1, "name": "Beta Feature", "key": "beta-feature", "is_simple_flag": False, "rollout_percentage": 100}
         ]
@@ -468,7 +544,7 @@ def test_feature_enabled_simple_with_none_rollout_percentage(self, patch_get):
     @mock.patch("posthog.client.decide")
     def test_feature_enabled_doesnt_exist(self, patch_decide, patch_poll):
         patch_decide.return_value = {"featureFlags": {}}
-        client = Client(TEST_API_KEY, personal_api_key="test")
+        client = Client(TEST_API_KEY)
         client.feature_flags = []
 
         self.assertFalse(client.feature_enabled("doesnt-exist", "distinct_id"))
@@ -477,7 +553,7 @@ def test_feature_enabled_doesnt_exist(self, patch_decide, patch_poll):
     @mock.patch("posthog.client.Poller")
     @mock.patch("posthog.client.decide")
     def test_personal_api_key_doesnt_exist(self, patch_decide, patch_poll):
-        client = Client(TEST_API_KEY)
+        client = Client(TEST_API_KEY, personal_api_key="test")
         client.feature_flags = []
 
         patch_decide.return_value = {"featureFlags": {"feature-flag": True}}