diff --git a/src/amplitude_experiment/assignment/assignment_config.py b/src/amplitude_experiment/assignment/assignment_config.py index 08ce2f8..db90c2c 100644 --- a/src/amplitude_experiment/assignment/assignment_config.py +++ b/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_user_props: bool = False, **kw): super(AssignmentConfig, self).__init__(**kw) self.cache_capacity = cache_capacity + self.send_evaluated_user_props = send_evaluated_user_props diff --git a/src/amplitude_experiment/assignment/assignment_service.py b/src/amplitude_experiment/assignment/assignment_service.py index 6ba48de..e08d975 100644 --- a/src/amplitude_experiment/assignment/assignment_service.py +++ b/src/amplitude_experiment/assignment/assignment_service.py @@ -8,9 +8,13 @@ FLAG_TYPE_HOLDOUT_GROUP = "holdout-group" -def to_event(assignment: Assignment) -> BaseEvent: +def to_event(assignment: Assignment, send_evaluated_user_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_user_props: + event.user_properties = assignment.user.user_properties + set_props = {} unset_props = {} @@ -46,10 +50,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_user_props: bool): self.amplitude = amplitude self.assignmentFilter = assignment_filter + self.send_evaluated_user_props = send_evaluated_user_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_user_props)) diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py index ba917d4..bc10ce9 100644 --- a/src/amplitude_experiment/local/client.py +++ b/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_user_props) self.logger = logging.getLogger("Amplitude") self.logger.addHandler(logging.StreamHandler()) if self.config.debug: diff --git a/tests/local/assignment/assignment_service_test.py b/tests/local/assignment/assignment_service_test.py index db95e44..5790461 100644 --- a/tests/local/assignment/assignment_service_test.py +++ b/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}) class AssignmentServiceTestCase(unittest.TestCase): @@ -61,7 +61,7 @@ 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) @@ -89,6 +89,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 +100,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() \ No newline at end of file + unittest.main()