Skip to content

Commit

Permalink
feat: add plan to config (#26)
Browse files Browse the repository at this point in the history
* add plan to config

* update plugin test to cover plan in context plugin
  • Loading branch information
bohan-amplitude committed Jun 2, 2022
1 parent f50042f commit eb3edf9
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 8 deletions.
14 changes: 9 additions & 5 deletions src/amplitude/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand Down
10 changes: 8 additions & 2 deletions src/amplitude/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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

Expand Down
7 changes: 6 additions & 1 deletion src/test/test_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand All @@ -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)
Expand Down

0 comments on commit eb3edf9

Please sign in to comment.