updated storage inmemory split storage

splitio/models/splits.py

@@ -42,7 +42,7 @@ def __init__(  # pylint: disable=too-many-arguments
             traffic_allocation=None,
             traffic_allocation_seed=None,
             configurations=None,
-            sets=[]
+            sets=None
     ):
         """
         Class constructor.

splitio/storage/inmemmory.py

@@ -23,6 +23,7 @@ def __init__(self):
         self._splits = {}
         self._change_number = -1
         self._traffic_types = Counter()
+        self._sets_feature_flag_map = {}
 
     def get(self, split_name):
         """
@@ -57,9 +58,15 @@ def put(self, split):
         """
         with self._lock:
             if split.name in self._splits:
+                self._remove_flag_sets(self._splits[split.name])
                 self._decrease_traffic_type_count(self._splits[split.name].traffic_type_name)
             self._splits[split.name] = split
             self._increase_traffic_type_count(split.traffic_type_name)
+            if split.sets is not None:
+                for flag_set in split.sets:
+                    if flag_set not in self._sets_feature_flag_map.keys():
+                        self._sets_feature_flag_map[flag_set] = set()
+                    self._sets_feature_flag_map[flag_set].add(split.name)
 
     def remove(self, split_name):
         """
@@ -79,11 +86,40 @@ def remove(self, split_name):
 
             self._splits.pop(split_name)
             self._decrease_traffic_type_count(split.traffic_type_name)
+            self._remove_flag_sets(split)
             return True
 
+    def _remove_flag_sets(self, feature_flag):
+        """
+        Remove flag sets associated to a split
+
+        :param feature_flag: feature flag object
+        :type feature_flag: splitio.models.splits.Split
+        """
+        if feature_flag.sets is not None:
+            for flag_set in feature_flag.sets:
+                self._sets_feature_flag_map[flag_set].remove(feature_flag.name)
+                if len(self._sets_feature_flag_map[flag_set]) == 0:
+                    del self._sets_feature_flag_map[flag_set]
+
+    def get_feature_flags_by_set(self, set):
+        """
+        Get list of feature flag names associated to a set, if it does not exist will return empty list
+
+        :param set: flag set
+        :type set: str
+
+        :return: list of feature flag names
+        :rtype: list
+        """
+        with self._lock:
+            if set not in self._sets_feature_flag_map:
+                return []
+            return list(self._sets_feature_flag_map[set])
+
     def get_change_number(self):
         """
-        Retrieve latest split change number.
+        Retrieve latest feature flag change number.
 
         :rtype: int
         """
@@ -102,27 +138,27 @@ def set_change_number(self, new_change_number):
 
     def get_split_names(self):
         """
-        Retrieve a list of all split names.
+        Retrieve a list of all feature flag names.
 
-        :return: List of split names.
+        :return: List of feature flag names.
         :rtype: list(str)
         """
         with self._lock:
             return list(self._splits.keys())
 
     def get_all_splits(self):
         """
-        Return all the splits.
+        Return all the feature flags.
 
-        :return: List of all the splits.
+        :return: List of all the feature flags.
         :rtype: list
         """
         with self._lock:
             return list(self._splits.values())
 
     def get_splits_count(self):
         """
-        Return splits count.
+        Return feature flags count.
 
         :rtype: int
         """
@@ -131,7 +167,7 @@ def get_splits_count(self):
 
     def is_valid_traffic_type(self, traffic_type_name):
         """
-        Return whether the traffic type exists in at least one split in cache.
+        Return whether the traffic type exists in at least one feature flag in cache.
 
         :param traffic_type_name: Traffic type to validate.
         :type traffic_type_name: str
@@ -142,12 +178,12 @@ def is_valid_traffic_type(self, traffic_type_name):
         with self._lock:
             return traffic_type_name in self._traffic_types
 
-    def kill_locally(self, split_name, default_treatment, change_number):
+    def kill_locally(self, feature_flag_name, default_treatment, change_number):
         """
-        Local kill for split
+        Local kill for feature flag
 
-        :param split_name: name of the split to perform kill
-        :type split_name: str
+        :param feature_flag_name: name of the feature flag to perform kill
+        :type feature_flag_name: str
         :param default_treatment: name of the default treatment to return
         :type default_treatment: str
         :param change_number: change_number
@@ -156,7 +192,7 @@ def kill_locally(self, split_name, default_treatment, change_number):
         with self._lock:
             if self.get_change_number() > change_number:
                 return
-            split = self._splits.get(split_name)
+            split = self._splits.get(feature_flag_name)
             if not split:
                 return
             split.local_kill(default_treatment, change_number)

tests/storage/test_inmemory_storage.py

@@ -9,7 +9,6 @@
 from splitio.models.events import Event, EventWrapper
 import splitio.models.telemetry as ModelTelemetry
 from splitio.engine.telemetry import TelemetryStorageProducer
-
 from splitio.storage.inmemmory import InMemorySplitStorage, InMemorySegmentStorage, \
     InMemoryImpressionStorage, InMemoryEventStorage, InMemoryTelemetryStorage
 
@@ -25,6 +24,9 @@ def test_storing_retrieving_splits(self, mocker):
         name_property = mocker.PropertyMock()
         name_property.return_value = 'some_split'
         type(split).name = name_property
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split).sets = sets_property
 
         storage.put(split)
         assert storage.get('some_split') == split
@@ -45,6 +47,10 @@ def test_get_splits(self, mocker):
         name2_prop = mocker.PropertyMock()
         name2_prop.return_value = 'split2'
         type(split2).name = name2_prop
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split1).sets = sets_property
+        type(split2).sets = sets_property
 
         storage = InMemorySplitStorage()
         storage.put(split1)
@@ -73,6 +79,10 @@ def test_get_split_names(self, mocker):
         name2_prop = mocker.PropertyMock()
         name2_prop.return_value = 'split2'
         type(split2).name = name2_prop
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split1).sets = sets_property
+        type(split2).sets = sets_property
 
         storage = InMemorySplitStorage()
         storage.put(split1)
@@ -90,6 +100,10 @@ def test_get_all_splits(self, mocker):
         name2_prop = mocker.PropertyMock()
         name2_prop.return_value = 'split2'
         type(split2).name = name2_prop
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split1).sets = sets_property
+        type(split2).sets = sets_property
 
         storage = InMemorySplitStorage()
         storage.put(split1)
@@ -120,6 +134,11 @@ def test_is_valid_traffic_type(self, mocker):
         type(split1).traffic_type_name = tt_user
         type(split2).traffic_type_name = tt_account
         type(split3).traffic_type_name = tt_user
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split1).sets = sets_property
+        type(split2).sets = sets_property
+        type(split3).sets = sets_property
 
         storage = InMemorySplitStorage()
 
@@ -155,11 +174,14 @@ def test_traffic_type_inc_dec_logic(self, mocker):
         name1_prop = mocker.PropertyMock()
         name1_prop.return_value = 'split1'
         type(split1).name = name1_prop
-
         split2 = mocker.Mock()
         name2_prop = mocker.PropertyMock()
         name2_prop.return_value = 'split1'
         type(split2).name = name2_prop
+        sets_property = mocker.PropertyMock()
+        sets_property.return_value = None
+        type(split1).sets = sets_property
+        type(split2).sets = sets_property
 
         tt_user = mocker.PropertyMock()
         tt_user.return_value = 'user'
@@ -198,6 +220,37 @@ def test_kill_locally(self):
         storage.kill_locally('some_split', 'default_treatment', 3)
         assert storage.get('some_split').change_number == 3
 
+    def test_flag_sets(self):
+        storage = InMemorySplitStorage()
+        assert storage._sets_feature_flag_map == {}
+
+        split1 = Split('split1', 123456789, False, 'some', 'traffic_type',
+                      'ACTIVE', 1, sets=['set10', 'set02'])
+        storage.put(split1)
+        assert storage.get_feature_flags_by_set('set10') == ['split1']
+        assert storage.get_feature_flags_by_set('set02') == ['split1']
+
+        split2 = Split('split2', 123456789, False, 'some', 'traffic_type',
+                      'ACTIVE', 1, sets=['set05', 'set02'])
+        storage.put(split2)
+        assert storage.get_feature_flags_by_set('set05') == ['split2']
+        assert sorted(storage.get_feature_flags_by_set('set02')) == ['split1', 'split2']
+
+        storage.remove(split2.name)
+        assert 'set5' not in storage._sets_feature_flag_map
+        assert storage.get_feature_flags_by_set('set02') == ['split1']
+        assert storage.get_feature_flags_by_set('set05') == []
+
+        split1 = Split('split1', 123456789, False, 'some', 'traffic_type',
+                      'ACTIVE', 1, sets=['set02'])
+        storage.put(split1)
+        assert 'set10' not in storage._sets_feature_flag_map
+        assert storage.get_feature_flags_by_set('set02') == ['split1']
+
+        storage.remove(split1.name)
+        assert storage._sets_feature_flag_map == {}
+        assert storage.get_feature_flags_by_set('set02') == []
+
 
 class InMemorySegmentStorageTests(object):
     """In memory segment storage tests."""