Added engine.impressions.adapters async classes

chillaq · chillaq · commit 71d17179edac · 2023-07-18T16:30:57.000-07:00
diff --git a/splitio/engine/impressions/adapters.py b/splitio/engine/impressions/adapters.py
@@ -21,7 +21,31 @@ def record_unique_keys(self, data):
         """
         pass
 
-class InMemorySenderAdapter(ImpressionsSenderAdapter):
+class InMemorySenderAdapterBase(ImpressionsSenderAdapter):
+    """In Memory Impressions Sender Adapter base class."""
+
+    def record_unique_keys(self, uniques):
+        """
+        post the unique keys to split back end.
+
+        :param uniques: unique keys disctionary
+        :type uniques: Dictionary {'feature_flag1': set(), 'feature_flag2': set(), .. }
+        """
+        pass
+
+    def _uniques_formatter(self, uniques):
+        """
+        Format the unique keys dictionary array to a JSON body
+
+        :param uniques: unique keys disctionary
+        :type uniques: Dictionary {'feature1_flag': set(), 'feature2_flag': set(), .. }
+
+        :return: unique keys JSON array
+        :rtype: json
+        """
+        return [{'f': feature, 'ks': list(keys)} for feature, keys in uniques.items()]
+
+class InMemorySenderAdapter(InMemorySenderAdapterBase):
     """In Memory Impressions Sender Adapter class."""
 
     def __init__(self, telemtry_http_client):
@@ -42,17 +66,28 @@ def record_unique_keys(self, uniques):
         """
         self._telemtry_http_client.record_unique_keys({'keys': self._uniques_formatter(uniques)})
 
-    def _uniques_formatter(self, uniques):
+
+class InMemorySenderAdapterAsync(InMemorySenderAdapterBase):
+    """In Memory Impressions Sender Adapter class."""
+
+    def __init__(self, telemtry_http_client):
         """
-        Format the unique keys dictionary array to a JSON body
+        Initialize In memory sender adapter instance
 
-        :param uniques: unique keys disctionary
-        :type uniques: Dictionary {'feature1_flag': set(), 'feature2_flag': set(), .. }
+        :param telemtry_http_client: instance of telemetry http api
+        :type telemtry_http_client: splitio.api.telemetry.TelemetryAPI
+        """
+        self._telemtry_http_client = telemtry_http_client
 
-        :return: unique keys JSON array
-        :rtype: json
+    async def record_unique_keys(self, uniques):
         """
-        return [{'f': feature, 'ks': list(keys)} for feature, keys in uniques.items()]
+        post the unique keys to split back end.
+
+        :param uniques: unique keys disctionary
+        :type uniques: Dictionary {'feature_flag1': set(), 'feature_flag2': set(), .. }
+        """
+        await self._telemtry_http_client.record_unique_keys({'keys': self._uniques_formatter(uniques)})
+
 
 class RedisSenderAdapter(ImpressionsSenderAdapter):
     """In Memory Impressions Sender Adapter class."""
@@ -118,6 +153,72 @@ 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)
 
+
+class RedisSenderAdapterAsync(ImpressionsSenderAdapter):
+    """In Memory Impressions Sender Adapter async class."""
+
+    def __init__(self, redis_client):
+        """
+        Initialize In memory sender adapter instance
+
+        :param telemtry_http_client: instance of telemetry http api
+        :type telemtry_http_client: splitio.api.telemetry.TelemetryAPI
+        """
+        self._redis_client = redis_client
+
+    async def record_unique_keys(self, uniques):
+        """
+        post the unique keys to redis.
+
+        :param uniques: unique keys disctionary
+        :type uniques: Dictionary {'feature_flag1': set(), 'feature_flag2': set(), .. }
+        """
+        bulk_mtks = _uniques_formatter(uniques)
+        try:
+            inserted = await self._redis_client.rpush(_MTK_QUEUE_KEY, *bulk_mtks)
+            await self._expire_keys(_MTK_QUEUE_KEY, _MTK_KEY_DEFAULT_TTL, inserted, len(bulk_mtks))
+            return True
+        except RedisAdapterException:
+            _LOGGER.error('Something went wrong when trying to add mtks to redis')
+            _LOGGER.error('Error: ', exc_info=True)
+            return False
+
+    async def flush_counters(self, to_send):
+        """
+        post the impression counters to redis.
+
+        :param to_send: unique keys disctionary
+        :type to_send: Dictionary {'feature_flag1': set(), 'feature_flag2': set(), .. }
+        """
+        try:
+            resulted = 0
+            counted = 0
+            pipe = self._redis_client.pipeline()
+            for pf_count in to_send:
+                pipe.hincrby(_IMP_COUNT_QUEUE_KEY, pf_count.feature + "::" + str(pf_count.timeframe), pf_count.count)
+                counted += pf_count.count
+            resulted = sum(await pipe.execute())
+            await self._expire_keys(_IMP_COUNT_QUEUE_KEY,
+                              _IMP_COUNT_KEY_DEFAULT_TTL, resulted, counted)
+            return True
+        except RedisAdapterException:
+            _LOGGER.error('Something went wrong when trying to add counters to redis')
+            _LOGGER.error('Error: ', exc_info=True)
+            return False
+
+    async 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:
+            await self._redis_client.expire(queue_key, key_default_ttl)
+
+
 class PluggableSenderAdapter(ImpressionsSenderAdapter):
     """In Memory Impressions Sender Adapter class."""
 
diff --git a/tests/engine/test_send_adapters.py b/tests/engine/test_send_adapters.py
@@ -2,11 +2,12 @@
 import ast
 import json
 import pytest
+import redis.asyncio as aioredis
 
-from splitio.engine.impressions.adapters import InMemorySenderAdapter, RedisSenderAdapter, PluggableSenderAdapter
+from splitio.engine.impressions.adapters import InMemorySenderAdapter, RedisSenderAdapter, PluggableSenderAdapter, InMemorySenderAdapterAsync, RedisSenderAdapterAsync
 from splitio.engine.impressions import adapters
-from splitio.api.telemetry import TelemetryAPI
-from splitio.storage.adapters.redis import RedisAdapter
+from splitio.api.telemetry import TelemetryAPI, TelemetryAPIAsync
+from splitio.storage.adapters.redis import RedisAdapter, RedisAdapterAsync
 from splitio.engine.impressions.manager import Counter
 from tests.storage.test_pluggable import StorageMockAdapter
 
@@ -43,6 +44,28 @@ def test_record_unique_keys(self, mocker):
 
         assert(mocker.called)
 
+
+class InMemorySenderAdapterAsyncTests(object):
+    """In memory sender adapter test."""
+
+    @pytest.mark.asyncio
+    async def test_record_unique_keys(self, mocker):
+        """Test sending unique keys."""
+
+        uniques = {"feature1": set({'key1', 'key2', 'key3'}),
+                   "feature2": set({'key1', 'key2', 'key3'}),
+                   }
+        telemetry_api = TelemetryAPIAsync(mocker.Mock(), 'some_api_key', mocker.Mock(), mocker.Mock())
+        self.called = False
+        async def record_unique_keys(*args):
+            self.called = True
+
+        telemetry_api.record_unique_keys = record_unique_keys
+        sender_adapter = InMemorySenderAdapterAsync(telemetry_api)
+        await sender_adapter.record_unique_keys(uniques)
+        assert(self.called)
+
+
 class RedisSenderAdapterTests(object):
     """Redis sender adapter test."""
 
@@ -103,6 +126,74 @@ def test_expire_keys(self, mocker):
         sender_adapter._expire_keys(mocker.Mock(), mocker.Mock(), total_keys, inserted)
         assert(mocker.called)
 
+
+class RedisSenderAdapterAsyncTests(object):
+    """Redis sender adapter test."""
+
+    @pytest.mark.asyncio
+    async def test_record_unique_keys(self, mocker):
+        """Test sending unique keys."""
+
+        uniques = {"feature1": set({'key1', 'key2', 'key3'}),
+                   "feature2": set({'key1', 'key2', 'key3'}),
+                   }
+        redis_client = RedisAdapterAsync(mocker.Mock(), mocker.Mock())
+        sender_adapter = RedisSenderAdapterAsync(redis_client)
+
+        self.called = False
+        async def rpush(*args):
+            self.called = True
+
+        redis_client.rpush = rpush
+        await sender_adapter.record_unique_keys(uniques)
+        assert(self.called)
+
+    @pytest.mark.asyncio
+    async def test_flush_counters(self, mocker):
+        """Test sending counters."""
+
+        counters = [
+            Counter.CountPerFeature('f1', 123, 2),
+            Counter.CountPerFeature('f2', 123, 123),
+        ]
+        redis_client = await aioredis.from_url("redis://localhost")
+        sender_adapter = RedisSenderAdapterAsync(redis_client)
+        self.called = False
+        def hincrby(*args):
+            self.called = True
+        self.called2 = False
+        async def execute(*args):
+            self.called2 = True
+            return [1]
+
+        with mock.patch('redis.asyncio.client.Pipeline.hincrby', hincrby):
+            with mock.patch('redis.asyncio.client.Pipeline.execute', execute):
+                await sender_adapter.flush_counters(counters)
+                assert(self.called)
+                assert(self.called2)
+
+    @pytest.mark.asyncio
+    async def test_expire_keys(self, mocker):
+        """Test set expire key."""
+
+        total_keys = 100
+        inserted = 10
+        redis_client = RedisAdapterAsync(mocker.Mock(), mocker.Mock())
+        sender_adapter = RedisSenderAdapterAsync(redis_client)
+        self.called = False
+        async def expire(*args):
+            self.called = True
+        redis_client.expire = expire
+
+        await sender_adapter._expire_keys(mocker.Mock(), mocker.Mock(), total_keys, inserted)
+        assert(not self.called)
+
+        total_keys = 100
+        inserted = 100
+        await sender_adapter._expire_keys(mocker.Mock(), mocker.Mock(), total_keys, inserted)
+        assert(self.called)
+
+
 class PluggableSenderAdapterTests(object):
     """Pluggable sender adapter test."""
 
