polishing

Author: chillaq

splitio/models/flag_sets.py

@@ -14,6 +14,89 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+class FlagSets(object):
+    """InMemory Flagsets storage."""
+
+    def __init__(self, flag_sets=[]):
+        """Constructor."""
+        self._lock = threading.RLock()
+        self.sets_feature_flag_map = {}
+        for flag_set in flag_sets:
+            self.sets_feature_flag_map[flag_set] = set()
+
+    def flag_set_exist(self, flag_set):
+        """
+        Check if a flagset exist in stored flagset
+
+        :param flag_set: set name
+        :type flag_set: str
+
+        :rtype: bool
+        """
+        with self._lock:
+            return flag_set in self.sets_feature_flag_map.keys()
+
+    def get_flag_set(self, flag_set):
+        """
+        fetch feature flags stored in a flag set
+
+        :param flag_set: set name
+        :type flag_set: str
+
+        :rtype: list(str)
+        """
+        with self._lock:
+            return self.sets_feature_flag_map.get(flag_set)
+
+    def add_flag_set(self, flag_set):
+        """
+        Add new flag set to storage
+
+        :param flag_set: set name
+        :type flag_set: str
+        """
+        with self._lock:
+            if not self.flag_set_exist(flag_set):
+                self.sets_feature_flag_map[flag_set] = set()
+
+    def remove_flag_set(self, flag_set):
+        """
+        Remove existing flag set from storage
+
+        :param flag_set: set name
+        :type flag_set: str
+        """
+        with self._lock:
+            if self.flag_set_exist(flag_set):
+                del self.sets_feature_flag_map[flag_set]
+
+    def add_feature_flag_to_flag_set(self, flag_set, feature_flag):
+        """
+        Add a feature flag to existing flag set
+
+        :param flag_set: set name
+        :type flag_set: str
+        :param feature_flag: feature flag name
+        :type feature_flag: str
+        """
+        with self._lock:
+            if self.flag_set_exist(flag_set):
+                self.sets_feature_flag_map[flag_set].add(feature_flag)
+
+    def remove_feature_flag_to_flag_set(self, flag_set, feature_flag):
+        """
+        Remove a feature flag from existing flag set
+
+        :param flag_set: set name
+        :type flag_set: str
+        :param feature_flag: feature flag name
+        :type feature_flag: str
+        """
+        with self._lock:
+            if self.flag_set_exist(flag_set):
+                self.sets_feature_flag_map[flag_set].remove(feature_flag)
+
+
 class InMemorySplitStorage(SplitStorage):
     """InMemory implementation of a feature flag storage."""
 

splitio/storage/inmemmory.py

@@ -1,7 +1,49 @@
 from splitio.storage import FlagSetsFilter
+from splitio.storage.inmemmory import FlagSets
 
 class FlagSetsFilterTests(object):
     """Flag sets filter storage tests."""
+    def test_without_initial_set(self):
+        flag_set = FlagSets()
+        assert flag_set.sets_feature_flag_map == {}
+
+        flag_set.add_flag_set('set1')
+        assert flag_set.get_flag_set('set1') == set({})
+        assert flag_set.flag_set_exist('set1') == True
+        assert flag_set.flag_set_exist('set2') == False
+
+        flag_set.add_feature_flag_to_flag_set('set1', 'split1')
+        assert flag_set.get_flag_set('set1') == {'split1'}
+        flag_set.add_feature_flag_to_flag_set('set1', 'split2')
+        assert flag_set.get_flag_set('set1') == {'split1', 'split2'}
+        flag_set.remove_feature_flag_to_flag_set('set1', 'split1')
+        assert flag_set.get_flag_set('set1') == {'split2'}
+        flag_set.remove_flag_set('set2')
+        assert flag_set.sets_feature_flag_map == {'set1': set({'split2'})}
+        flag_set.remove_flag_set('set1')
+        assert flag_set.sets_feature_flag_map == {}
+        assert flag_set.flag_set_exist('set1') == False
+
+    def test_with_initial_set(self):
+        flag_set = FlagSets(['set1', 'set2'])
+        assert flag_set.sets_feature_flag_map == {'set1': set(), 'set2': set()}
+
+        flag_set.add_flag_set('set1')
+        assert flag_set.get_flag_set('set1') == set({})
+        assert flag_set.flag_set_exist('set1') == True
+        assert flag_set.flag_set_exist('set2') == True
+
+        flag_set.add_feature_flag_to_flag_set('set1', 'split1')
+        assert flag_set.get_flag_set('set1') == {'split1'}
+        flag_set.add_feature_flag_to_flag_set('set1', 'split2')
+        assert flag_set.get_flag_set('set1') == {'split1', 'split2'}
+        flag_set.remove_feature_flag_to_flag_set('set1', 'split1')
+        assert flag_set.get_flag_set('set1') == {'split2'}
+        flag_set.remove_flag_set('set2')
+        assert flag_set.sets_feature_flag_map == {'set1': set({'split2'})}
+        flag_set.remove_flag_set('set1')
+        assert flag_set.sets_feature_flag_map == {}
+        assert flag_set.flag_set_exist('set1') == False
 
     def test_flag_set_filter(self):
         flag_set_filter = FlagSetsFilter()