feat: support sending evaluated user properties in local evaluation assignment tracking (#51)

Author: tyiuhc

src/amplitude_experiment/assignment/assignment_config.py

@@ -2,6 +2,7 @@
 
 
 class AssignmentConfig(amplitude.Config):
-    def __init__(self, cache_capacity: int = 65536, **kw):
+    def __init__(self, cache_capacity: int = 65536, send_evaluated_props: bool = False, **kw):
         super(AssignmentConfig, self).__init__(**kw)
         self.cache_capacity = cache_capacity
+        self.send_evaluated_props = send_evaluated_props

src/amplitude_experiment/assignment/assignment_service.py

@@ -8,9 +8,21 @@
 FLAG_TYPE_HOLDOUT_GROUP = "holdout-group"
 
 
-def to_event(assignment: Assignment) -> BaseEvent:
+def to_event(assignment: Assignment, send_evaluated_props: bool) -> BaseEvent:
     event = BaseEvent(event_type='[Experiment] Assignment', user_id=assignment.user.user_id,
                       device_id=assignment.user.device_id, event_properties={}, user_properties={})
+
+    # If send_evaluated_props is True, populate event with all relevant user attributes
+    if send_evaluated_props:
+        user_attributes = [
+            "country", "city", "region", "dma", "language",
+            "platform", "version", "os", "device_manufacturer",
+            "device_brand", "device_model", "carrier"
+        ]
+        for attr in user_attributes:
+            setattr(event, attr, getattr(assignment.user, attr, None))
+        event.user_properties = assignment.user.user_properties
+
     set_props = {}
     unset_props = {}
 
@@ -46,10 +58,11 @@ def to_event(assignment: Assignment) -> BaseEvent:
 
 
 class AssignmentService:
-    def __init__(self, amplitude: Amplitude, assignment_filter: AssignmentFilter):
+    def __init__(self, amplitude: Amplitude, assignment_filter: AssignmentFilter, send_evaluated_props: bool):
         self.amplitude = amplitude
         self.assignmentFilter = assignment_filter
+        self.send_evaluated_props = send_evaluated_props
 
     def track(self, assignment: Assignment):
         if self.assignmentFilter.should_track(assignment):
-            self.amplitude.track(to_event(assignment))
+            self.amplitude.track(to_event(assignment, self.send_evaluated_props))

src/amplitude_experiment/local/client.py

@@ -47,7 +47,7 @@ def __init__(self, api_key: str, config: LocalEvaluationConfig = None):
         if config and config.assignment_config:
             instance = Amplitude(config.assignment_config.api_key, config.assignment_config)
             self.assignment_service = AssignmentService(instance, AssignmentFilter(
-                config.assignment_config.cache_capacity))
+                config.assignment_config.cache_capacity), config.assignment_config.send_evaluated_props)
         self.logger = logging.getLogger("Amplitude")
         self.logger.addHandler(logging.StreamHandler())
         if self.config.debug:

tests/local/assignment/assignment_service_test.py

@@ -8,7 +8,7 @@
 from src.amplitude_experiment.assignment import AssignmentFilter, Assignment, DAY_MILLIS, to_event, AssignmentService
 from src.amplitude_experiment.util import hash_code
 
-user = User(user_id='user', device_id='device')
+user = User(user_id='user', device_id='device', user_properties={'user_prop': True}, country='country')
 
 
 class AssignmentServiceTestCase(unittest.TestCase):
@@ -61,10 +61,11 @@ def test_to_event(self):
             'empty_variant': empty_variant,
         }
         assignment = Assignment(user, results)
-        event = to_event(assignment)
+        event = to_event(assignment, True)
         self.assertEqual(user.user_id, event.user_id)
         self.assertEqual(user.device_id, event.device_id)
         self.assertEqual('[Experiment] Assignment', event.event_type)
+        self.assertEqual('country', event.country)
         # Validate event properties
         event_properties = event.event_properties
         self.assertEqual('control', event_properties['basic.variant'])
@@ -89,6 +90,7 @@ def test_to_event(self):
         self.assertEqual('on', set_properties['[Experiment] empty_metadata'])
         unset_properties = user_properties['$unset']
         self.assertEqual('-', unset_properties['[Experiment] default'])
+        self.assertTrue(user_properties['user_prop'])
 
         # Validate insert id
         canonicalization = 'user device basic control default off different_value on empty_metadata on holdout ' \
@@ -99,11 +101,11 @@ def test_to_event(self):
     def test_tracking_called(self):
         instance = Amplitude('')
         instance.track = MagicMock()
-        service = AssignmentService(instance, AssignmentFilter(2))
+        service = AssignmentService(instance, AssignmentFilter(2), False)
         results = {'flag-key-1': Variant(key='on')}
         service.track(Assignment(user, results))
         self.assertTrue(instance.track.called)
 
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()