feat: add plan to config (#26)

* add plan to config

* update plugin test to cover plan in context plugin

Author: bohan-amplitude

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.

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
 

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)