Merge pull request #331 from splitio/pluggable-mtk-adapter

Added Pluggable support for MTK to sender adapter

Author: chillaq

splitio/engine/impressions/adapters.py

@@ -74,7 +74,7 @@ def record_unique_keys(self, uniques):
         :param uniques: unique keys disctionary
         :type uniques: Dictionary {'feature1': set(), 'feature2': set(), .. }
         """
-        bulk_mtks = self._uniques_formatter(uniques)
+        bulk_mtks = _uniques_formatter(uniques)
         try:
             inserted = self._redis_client.rpush(self.MTK_QUEUE_KEY, *bulk_mtks)
             self._expire_keys(self.MTK_QUEUE_KEY, self.MTK_KEY_DEFAULT_TTL, inserted, len(bulk_mtks))
@@ -119,14 +119,82 @@ def _expire_keys(self, queue_key, key_default_ttl, total_keys, inserted):
         if total_keys == inserted:
             self._redis_client.expire(queue_key, key_default_ttl)
 
-    def _uniques_formatter(self, uniques):
+class PluggableSenderAdapter(ImpressionsSenderAdapter):
+    """In Memory Impressions Sender Adapter class."""
+
+    MTK_QUEUE_KEY = 'SPLITIO.uniquekeys'
+    MTK_KEY_DEFAULT_TTL = 3600
+    IMP_COUNT_QUEUE_KEY = 'SPLITIO.impressions.count'
+    IMP_COUNT_KEY_DEFAULT_TTL = 3600
+
+    def __init__(self, adapter_client, prefix=None):
         """
-        Format the unique keys dictionary array to a JSON body
+        Initialize pluggable sender adapter instance
+
+        :param telemtry_http_client: instance of telemetry http api
+        :type telemtry_http_client: splitio.api.telemetry.TelemetryAPI
+        """
+        self._adapter_client = adapter_client
+        self._prefix = ""
+        if prefix is not None:
+            self._prefix = prefix + "."
+
+    def record_unique_keys(self, uniques):
+        """
+        post the unique keys to storage.
 
         :param uniques: unique keys disctionary
         :type uniques: Dictionary {'feature1': set(), 'feature2': set(), .. }
+        """
+        bulk_mtks = _uniques_formatter(uniques)
+        try:
+            inserted = self._adapter_client.push_items(self.MTK_QUEUE_KEY, *bulk_mtks)
+            self._expire_keys(self._prefix + self.MTK_QUEUE_KEY, self.MTK_KEY_DEFAULT_TTL, inserted, len(bulk_mtks))
+            return True
+        except RedisAdapterException:
+            _LOGGER.error('Something went wrong when trying to add mtks to storage adapter')
+            _LOGGER.error('Error: ', exc_info=True)
+            return False
 
-        :return: unique keys JSON array
-        :rtype: json
+    def flush_counters(self, to_send):
         """
-        return [json.dumps({'f': feature, 'ks': list(keys)}) for feature, keys in uniques.items()]
+        post the impression counters to storage.
+
+        :param to_send: unique keys disctionary
+        :type to_send: Dictionary {'feature1': set(), 'feature2': set(), .. }
+        """
+        try:
+            resulted = 0
+            for pf_count in to_send:
+                resulted = self._adapter_client.increment(self._prefix + self.IMP_COUNT_QUEUE_KEY + "." + pf_count.feature + "::" + str(pf_count.timeframe), pf_count.count)
+                self._expire_keys(self._prefix + self.IMP_COUNT_QUEUE_KEY + "." + pf_count.feature + "::" + str(pf_count.timeframe),
+                              self.IMP_COUNT_KEY_DEFAULT_TTL, resulted, pf_count.count)
+            return True
+        except RedisAdapterException:
+            _LOGGER.error('Something went wrong when trying to add counters to storage adapter')
+            _LOGGER.error('Error: ', exc_info=True)
+            return False
+
+    def _expire_keys(self, queue_key, key_default_ttl, total_keys, inserted):
+        """
+        Set expire
+
+        :param total_keys: length of keys.
+        :type total_keys: int
+        :param inserted: added keys.
+        :type inserted: int
+        """
+        if total_keys == inserted:
+            self._adapter_client.expire(queue_key, key_default_ttl)
+
+def _uniques_formatter(uniques):
+    """
+    Format the unique keys dictionary array to a JSON body
+
+    :param uniques: unique keys disctionary
+    :type uniques: Dictionary {'feature1': set(), 'feature2': set(), .. }
+
+    :return: unique keys JSON array
+    :rtype: json
+    """
+    return [json.dumps({'f': feature, 'ks': list(keys)}) for feature, keys in uniques.items()]

tests/engine/test_send_adapters.py

@@ -1,10 +1,15 @@
 import unittest.mock as mock
 import ast
+import json
+import pytest
 
-from splitio.engine.impressions.adapters import InMemorySenderAdapter, RedisSenderAdapter
+from splitio.engine.impressions.adapters import InMemorySenderAdapter, RedisSenderAdapter, PluggableSenderAdapter
+from splitio.engine.impressions import adapters
 from splitio.api.telemetry import TelemetryAPI
 from splitio.storage.adapters.redis import RedisAdapter
 from splitio.engine.impressions.manager import Counter
+from tests.storage.test_pluggable import StorageMockAdapter
+
 
 class InMemorySenderAdapterTests(object):
     """In memory sender adapter test."""
@@ -52,9 +57,8 @@ def test_uniques_formatter(self, mocker):
             {'f': 'feature2', 'ks': ['key6', 'key1', 'key10']},
         ]
 
-        sender_adapter = RedisSenderAdapter(mocker.Mock())
         for i in range(0,1):
-            assert(sorted(ast.literal_eval(sender_adapter._uniques_formatter(uniques)[i])["ks"]) == sorted(formatted[i]["ks"]))
+            assert(sorted(ast.literal_eval(adapters._uniques_formatter(uniques)[i])["ks"]) == sorted(formatted[i]["ks"]))
 
     @mock.patch('splitio.storage.adapters.redis.RedisAdapter.rpush')
     def test_record_unique_keys(self, mocker):
@@ -98,3 +102,45 @@ def test_expire_keys(self, mocker):
         inserted = 100
         sender_adapter._expire_keys(mocker.Mock(), mocker.Mock(), total_keys, inserted)
         assert(mocker.called)
+
+class PluggableSenderAdapterTests(object):
+    """Pluggable sender adapter test."""
+
+    def test_record_unique_keys(self, mocker):
+        """Test sending unique keys."""
+        adapter = StorageMockAdapter()
+        sender_adapter = PluggableSenderAdapter(adapter)
+
+        uniques = {"feature1": set({"key1", "key2", "key3"}),
+                   "feature2": set({"key1", "key6", "key10"}),
+                   }
+        formatted = [
+            '{"f": "feature1", "ks": ["key3", "key2", "key1"]}',
+            '{"f": "feature2", "ks": ["key1", "key10", "key6"]}',
+        ]
+
+        sender_adapter.record_unique_keys(uniques)
+        assert(sorted(json.loads(adapter._keys[sender_adapter.MTK_QUEUE_KEY][0])["ks"]) == sorted(json.loads(formatted[0])["ks"]))
+        assert(sorted(json.loads(adapter._keys[sender_adapter.MTK_QUEUE_KEY][1])["ks"]) == sorted(json.loads(formatted[1])["ks"]))
+        assert(json.loads(adapter._keys[sender_adapter.MTK_QUEUE_KEY][0])["f"] == "feature1")
+        assert(json.loads(adapter._keys[sender_adapter.MTK_QUEUE_KEY][1])["f"] == "feature2")
+        assert(adapter._expire[sender_adapter.MTK_QUEUE_KEY] == sender_adapter.MTK_KEY_DEFAULT_TTL)
+        sender_adapter.record_unique_keys(uniques)
+        assert(adapter._expire[sender_adapter.MTK_QUEUE_KEY] != -1)
+
+    def test_flush_counters(self, mocker):
+        """Test sending counters."""
+        adapter = StorageMockAdapter()
+        sender_adapter = PluggableSenderAdapter(adapter)
+
+        counters = [
+            Counter.CountPerFeature('f1', 123, 2),
+            Counter.CountPerFeature('f2', 123, 123),
+        ]
+
+        sender_adapter.flush_counters(counters)
+        assert(adapter._keys[sender_adapter.IMP_COUNT_QUEUE_KEY + "." + 'f1::123'] == 2)
+        assert(adapter._keys[sender_adapter.IMP_COUNT_QUEUE_KEY + "." + 'f2::123'] == 123)
+        assert(adapter._expire[sender_adapter.IMP_COUNT_QUEUE_KEY + "." + 'f1::123'] == sender_adapter.IMP_COUNT_KEY_DEFAULT_TTL)
+        sender_adapter.flush_counters(counters)
+        assert(adapter._expire[sender_adapter.IMP_COUNT_QUEUE_KEY + "." + 'f2::123'] == sender_adapter.IMP_COUNT_KEY_DEFAULT_TTL)