From eb3edf969a29e372da054fb44d2ea8cc89b09d42 Mon Sep 17 00:00:00 2001 From: bohan-amplitude <96497806+bohan-amplitude@users.noreply.github.com> Date: Thu, 2 Jun 2022 11:03:45 -0700 Subject: [PATCH] feat: add plan to config (#26) * add plan to config * update plugin test to cover plan in context plugin --- src/amplitude/config.py | 14 +++++++++----- src/amplitude/plugin.py | 10 ++++++++-- src/test/test_plugin.py | 7 ++++++- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/src/amplitude/config.py b/src/amplitude/config.py index dfc16eb..ac70308 100644 --- a/src/amplitude/config.py +++ b/src/amplitude/config.py @@ -8,7 +8,7 @@ from typing import Optional, Callable from amplitude import constants -from amplitude.event import BaseEvent +from amplitude.event import BaseEvent, Plan from amplitude.storage import InMemoryStorageProvider, StorageProvider, Storage @@ -21,6 +21,7 @@ class Config: Must be positive. flush_interval_millis (int, optional): The events buffered in memory will wait no longer than flush_interval_millis. Must be positive. + flush_max_retries (int, optional): The maximum retry attempts for an event when receiving error response. logger (optional): The logger used by Amplitude client. Default to logging.getLogger(constants.LOGGER_NAME). min_id_length (int, optional): The minimum length of user_id and device_id for events. Default to 5. callback (callable, optional): The client level callback function. Triggered on every events sent or failed. @@ -29,8 +30,9 @@ class Config: use_batch(bool, optional): True to use batch API endpoint, False to use HTTP V2 API endpoint. Default to False. server_url (str, optional): API endpoint url. Default to None. Auto selected by configured server_zone and use_batch if set to None. Support customized url by setting string value. - storage_provider(amplitude.storage.StorageProvider, optional): Default to InMemoryStorageProvider. + storage_provider (amplitude.storage.StorageProvider, optional): Default to InMemoryStorageProvider. Provide storage instance for events buffer. + plan (amplitude.event.Plan, optional): Tracking plan information. Default to None. Properties: options: A dictionary contains minimum id length information. None if min_id_length not set. @@ -52,9 +54,10 @@ def __init__(self, api_key: str = None, server_zone: str = constants.DEFAULT_ZONE, use_batch: bool = False, server_url: Optional[str] = None, - storage_provider: StorageProvider = InMemoryStorageProvider()): + storage_provider: StorageProvider = InMemoryStorageProvider(), + plan: Plan = None): """The constructor of Config class""" - self.api_key = api_key + self.api_key: str = api_key self._flush_queue_size: int = flush_queue_size self._flush_size_divider: int = 1 self.flush_interval_millis: int = flush_interval_millis @@ -66,7 +69,8 @@ def __init__(self, api_key: str = None, self.use_batch: bool = use_batch self._url: Optional[str] = server_url self.storage_provider: StorageProvider = storage_provider - self.opt_out = False + self.opt_out: bool = False + self.plan: Plan = plan def get_storage(self) -> Storage: """Use configured StorageProvider to create a Storage instance then return. diff --git a/src/amplitude/plugin.py b/src/amplitude/plugin.py index 7c9c73b..1956cc3 100644 --- a/src/amplitude/plugin.py +++ b/src/amplitude/plugin.py @@ -187,6 +187,10 @@ def __init__(self): """The constructor of ContextPlugin class""" super().__init__(constants.PluginType.BEFORE) self.context_string = f"{constants.SDK_LIBRARY}/{constants.SDK_VERSION}" + self.configuration = None + + def setup(self, client): + self.configuration = client.configuration def apply_context_data(self, event: BaseEvent): """Add SDK name and version to event.library. @@ -203,9 +207,11 @@ def execute(self, event: BaseEvent) -> BaseEvent: event (BaseEvent): The event to be processed. """ if not event.time: - event.time = utils.current_milliseconds() + event["time"] = utils.current_milliseconds() if not event.insert_id: - event.insert_id = str(uuid.uuid4()) + event["insert_id"] = str(uuid.uuid4()) + if self.configuration.plan and (not event.plan): + event["plan"] = self.configuration.plan self.apply_context_data(event) return event diff --git a/src/test/test_plugin.py b/src/test/test_plugin.py index a1d8219..ee342f7 100644 --- a/src/test/test_plugin.py +++ b/src/test/test_plugin.py @@ -2,7 +2,7 @@ from unittest.mock import MagicMock from amplitude.plugin import AmplitudeDestinationPlugin, ContextPlugin, EventPlugin, DestinationPlugin -from amplitude import Amplitude, PluginType, BaseEvent, RevenueEvent, IdentifyEvent, GroupIdentifyEvent, Config +from amplitude import Amplitude, PluginType, BaseEvent, RevenueEvent, IdentifyEvent, GroupIdentifyEvent, Config, Plan class AmplitudePluginTestCase(unittest.TestCase): @@ -25,14 +25,18 @@ def test_plugin_initialize_amplitude_client_context_plugin_creation_success(self def test_plugin_context_plugin_execute_event_success(self): context_plugin = ContextPlugin() + context_plugin.setup(Amplitude("test_api_key")) + context_plugin.configuration.plan = Plan(source="test_source") event = BaseEvent("test_event", user_id="test_user") self.assertIsNone(event.time) self.assertIsNone(event.insert_id) self.assertIsNone(event.library) + self.assertIsNone(event.plan) context_plugin.execute(event) self.assertTrue(isinstance(event.time, int)) self.assertTrue(isinstance(event.insert_id, str)) self.assertTrue(isinstance(event.library, str)) + self.assertTrue(isinstance(event.plan, Plan)) def test_plugin_event_plugin_process_event_success(self): plugin = EventPlugin(PluginType.BEFORE) @@ -50,6 +54,7 @@ def test_plugin_destination_plugin_add_remove_plugin_success(self): destination_plugin = DestinationPlugin() destination_plugin.timeline.configuration = Config() context_plugin = ContextPlugin() + context_plugin.configuration = destination_plugin.timeline.configuration event = BaseEvent("test_event", user_id="test_user") destination_plugin.add(context_plugin) destination_plugin.execute(event)