Merge pull request #402 from splitio/async-memory-segment-storage

added memory segment storage async class

Author: chillaq

splitio/storage/inmemmory.py

@@ -576,6 +576,134 @@ def get_segments_keys_count(self):
             return total_count
 
 
+class InMemorySegmentStorageAsync(SegmentStorage):
+    """In-memory implementation of a segment async storage."""
+
+    def __init__(self):
+        """Constructor."""
+        self._segments = {}
+        self._change_numbers = {}
+        self._lock = asyncio.Lock()
+
+    async def get(self, segment_name):
+        """
+        Retrieve a segment.
+
+        :param segment_name: Name of the segment to fetch.
+        :type segment_name: str
+
+        :rtype: str
+        """
+        async with self._lock:
+            fetched = self._segments.get(segment_name)
+            if fetched is None:
+                _LOGGER.debug(
+                    "Tried to retrieve nonexistant segment %s. Skipping",
+                    segment_name
+                )
+            return fetched
+
+    async def put(self, segment):
+        """
+        Store a segment.
+
+        :param segment: Segment to store.
+        :type segment: splitio.models.segment.Segment
+        """
+        async with self._lock:
+            self._segments[segment.name] = segment
+
+    async def update(self, segment_name, to_add, to_remove, change_number=None):
+        """
+        Update a split. Create it if it doesn't exist.
+
+        :param segment_name: Name of the segment to update.
+        :type segment_name: str
+        :param to_add: Set of members to add to the segment.
+        :type to_add: set
+        :param to_remove: List of members to remove from the segment.
+        :type to_remove: Set
+        """
+        async with self._lock:
+            if segment_name not in self._segments:
+                self._segments[segment_name] = Segment(segment_name, to_add, change_number)
+                return
+
+            self._segments[segment_name].update(to_add, to_remove)
+            if change_number is not None:
+                self._segments[segment_name].change_number = change_number
+
+    async def get_change_number(self, segment_name):
+        """
+        Retrieve latest change number for a segment.
+
+        :param segment_name: Name of the segment.
+        :type segment_name: str
+
+        :rtype: int
+        """
+        async with self._lock:
+            if segment_name not in self._segments:
+                return None
+            return self._segments[segment_name].change_number
+
+    async def set_change_number(self, segment_name, new_change_number):
+        """
+        Set the latest change number.
+
+        :param segment_name: Name of the segment.
+        :type segment_name: str
+        :param new_change_number: New change number.
+        :type new_change_number: int
+        """
+        async with self._lock:
+            if segment_name not in self._segments:
+                return
+            self._segments[segment_name].change_number = new_change_number
+
+    async def segment_contains(self, segment_name, key):
+        """
+        Check whether a specific key belongs to a segment in storage.
+
+        :param segment_name: Name of the segment to search in.
+        :type segment_name: str
+        :param key: Key to search for.
+        :type key: str
+
+        :return: True if the segment contains the key. False otherwise.
+        :rtype: bool
+        """
+        async with self._lock:
+            if segment_name not in self._segments:
+                _LOGGER.warning(
+                    "Tried to query members for nonexistant segment %s. Returning False",
+                    segment_name
+                )
+                return False
+            return self._segments[segment_name].contains(key)
+
+    async def get_segments_count(self):
+        """
+        Retrieve segments count.
+
+        :rtype: int
+        """
+        async with self._lock:
+            return len(self._segments)
+
+    async def get_segments_keys_count(self):
+        """
+        Retrieve segments keys count.
+
+        :rtype: int
+        """
+        total_count = 0
+        async with self._lock:
+            for segment in self._segments:
+                total_count += len(self._segments[segment]._keys)
+            return total_count
+
+
 class InMemoryImpressionStorage(ImpressionStorage):
     """In memory implementation of an impressions storage."""
 

tests/storage/test_inmemory_storage.py

@@ -11,8 +11,7 @@
 from splitio.engine.telemetry import TelemetryStorageProducer
 
 from splitio.storage.inmemmory import InMemorySplitStorage, InMemorySegmentStorage, \
-    InMemoryImpressionStorage, InMemoryEventStorage, InMemoryTelemetryStorage, InMemorySplitStorageAsync
-
+    InMemoryImpressionStorage, InMemoryEventStorage, InMemoryTelemetryStorage, InMemorySegmentStorageAsync, InMemorySplitStorageAsync
 
 class InMemorySplitStorageTests(object):
     """In memory split storage test cases."""
@@ -455,6 +454,71 @@ def test_segment_update(self):
         assert storage.get_change_number('some_segment') == 456
 
 
+class InMemorySegmentStorageAsyncTests(object):
+    """In memory segment storage tests."""
+
+    @pytest.mark.asyncio
+    async def test_segment_storage_retrieval(self, mocker):
+        """Test storing and retrieving segments."""
+        storage = InMemorySegmentStorageAsync()
+        segment = mocker.Mock(spec=Segment)
+        name_property = mocker.PropertyMock()
+        name_property.return_value = 'some_segment'
+        type(segment).name = name_property
+
+        await storage.put(segment)
+        assert await storage.get('some_segment') == segment
+        assert await storage.get('nonexistant-segment') is None
+
+    @pytest.mark.asyncio
+    async def test_change_number(self, mocker):
+        """Test storing and retrieving segment changeNumber."""
+        storage = InMemorySegmentStorageAsync()
+        await storage.set_change_number('some_segment', 123)
+        # Change number is not updated if segment doesn't exist
+        assert await storage.get_change_number('some_segment') is None
+        assert await storage.get_change_number('nonexistant-segment') is None
+
+        # Change number is updated if segment does exist.
+        storage = InMemorySegmentStorageAsync()
+        segment = mocker.Mock(spec=Segment)
+        name_property = mocker.PropertyMock()
+        name_property.return_value = 'some_segment'
+        type(segment).name = name_property
+        await storage.put(segment)
+        await storage.set_change_number('some_segment', 123)
+        assert await storage.get_change_number('some_segment') == 123
+
+    @pytest.mark.asyncio
+    async def test_segment_contains(self, mocker):
+        """Test using storage to determine whether a key belongs to a segment."""
+        storage = InMemorySegmentStorageAsync()
+        segment = mocker.Mock(spec=Segment)
+        name_property = mocker.PropertyMock()
+        name_property.return_value = 'some_segment'
+        type(segment).name = name_property
+        await storage.put(segment)
+
+        await storage.segment_contains('some_segment', 'abc')
+        assert segment.contains.mock_calls[0] == mocker.call('abc')
+
+    @pytest.mark.asyncio
+    async def test_segment_update(self):
+        """Test updating a segment."""
+        storage = InMemorySegmentStorageAsync()
+        segment = Segment('some_segment', ['key1', 'key2', 'key3'], 123)
+        await storage.put(segment)
+        assert await storage.get('some_segment') == segment
+
+        await storage.update('some_segment', ['key4', 'key5'], ['key2', 'key3'], 456)
+        assert await storage.segment_contains('some_segment', 'key1')
+        assert await storage.segment_contains('some_segment', 'key4')
+        assert await storage.segment_contains('some_segment', 'key5')
+        assert not await storage.segment_contains('some_segment', 'key2')
+        assert not await storage.segment_contains('some_segment', 'key3')
+        assert await storage.get_change_number('some_segment') == 456
+
+
 class InMemoryImpressionsStorageTests(object):
     """InMemory impressions storage test cases."""