Updated redis storage to reflect flagset filter

chillaq · chillaq · commit 878b0a6da388 · 2023-10-30T13:55:32.000-07:00
diff --git a/splitio/storage/redis.py b/splitio/storage/redis.py
@@ -89,7 +89,14 @@ def get(self, feature_flag_name):  # pylint: disable=method-hidden
             raw = self._redis.get(self._get_key(feature_flag_name))
             _LOGGER.debug("Fetchting Feature flag [%s] from redis" % feature_flag_name)
             _LOGGER.debug(raw)
-            return splits.from_raw(json.loads(raw)) if raw is not None else None
+            if raw is None:
+                return None
+
+            feature_flag = splits.from_raw(json.loads(raw))
+            if self.flag_set_filter.intersect(feature_flag.sets):
+                return feature_flag
+
+            return None
         except RedisAdapterException:
             _LOGGER.error('Error fetching feature flag from storage')
             _LOGGER.debug('Error: ', exc_info=True)
@@ -206,6 +213,9 @@ def get_split_names(self):
         :return: List of feature flag names.
         :rtype: list(str)
         """
+        if self.flag_set_filter.should_filter:
+            return self.get_feature_flags_by_sets(self.flag_set_filter.flag_sets)
+
         try:
             keys = self._redis.keys(self._get_key('*'))
             _LOGGER.debug("Fetchting feature flag names from redis: %s" % keys)
@@ -229,7 +239,12 @@ def get_all_splits(self):
         :return: List of all feature flags in cache.
         :rtype: list(splitio.models.splits.Split)
         """
-        keys = self._redis.keys(self._get_key('*'))
+        if self.flag_set_filter.should_filter:
+            keys = self.get_feature_flags_by_sets(self.flag_set_filter.flag_sets)
+        else:
+            keys = self._redis.keys(self._get_key('*'))
+        if keys == []:
+            return []
         to_return = []
         try:
             _LOGGER.debug("Fetchting all feature flags from redis: %s" % keys)
diff --git a/tests/storage/test_redis.py b/tests/storage/test_redis.py
@@ -15,6 +15,8 @@
 from splitio.models.events import Event, EventWrapper
 from splitio.models.telemetry import MethodExceptions, MethodLatencies, TelemetryConfig, MethodExceptionsAndLatencies
 from splitio.storage import FlagSetsFilter
+from tests.integration import splits_json
+
 
 class RedisSplitStorageTests(object):
     """Redis split storage test cases."""
@@ -187,6 +189,52 @@ def test_flag_sets(self, mocker):
         storage2 = RedisSplitStorage(adapter, True, 1, ['set2', 'set3'])
         assert storage2.flag_set_filter.flag_sets == set({'set2', 'set3'})
 
+    def test_fetching_split_with_flag_set(self, mocker):
+        """Test retrieving a split works."""
+        adapter = mocker.Mock(spec=RedisAdapter)
+        adapter.get.return_value = json.dumps(splits_json["splitChange1_1"]["splits"][0])
+        adapter.keys.return_value = ['SPLIT_1', 'SPLIT_2']
+
+        def mget(keys):
+            if keys == ['SPLIT_2']:
+                return [json.dumps(splits_json["splitChange1_1"]["splits"][0])]
+            if keys == ['SPLIT_2', 'SPLIT_1']:
+                return [json.dumps(splits_json["splitChange1_1"]["splits"][0]), json.dumps(splits_json["splitChange1_1"]["splits"][1])]
+        adapter.mget = mget
+
+        storage = RedisSplitStorage(adapter, config_flag_sets=['set_1'])
+
+        def get_feature_flags_by_sets(flag_sets):
+            if flag_sets=={'set_1'}:
+                return []
+            if flag_sets=={'set2'}:
+                return ['SPLIT_2']
+            if flag_sets=={'set2', 'set1'}:
+                return ['SPLIT_2', 'SPLIT_1']
+        storage.get_feature_flags_by_sets = get_feature_flags_by_sets
+
+        assert storage.get('SPLIT_2') == None
+        assert storage.get_split_names() == []
+        assert storage.get_all_splits() == []
+
+        storage = RedisSplitStorage(adapter, config_flag_sets=['set2'])
+        storage.get_feature_flags_by_sets = get_feature_flags_by_sets
+        assert storage.get('SPLIT_2').name == 'SPLIT_2'
+        assert storage.get_split_names() == ['SPLIT_2']
+        splits = storage.get_all_splits()
+        assert splits[0].name == 'SPLIT_2'
+        assert len(splits) == 1
+
+        storage = RedisSplitStorage(adapter, config_flag_sets=['set2', 'set1'])
+        storage.get_feature_flags_by_sets = get_feature_flags_by_sets
+        assert storage.get('SPLIT_2').name == 'SPLIT_2'
+        assert storage.get_split_names() == ['SPLIT_2', 'SPLIT_1']
+        splits = storage.get_all_splits()
+        assert splits[0].name == 'SPLIT_2'
+        assert splits[1].name == 'SPLIT_1'
+        assert len(splits) == 2
+
+
 class RedisSegmentStorageTests(object):
     """Redis segment storage test cases."""
 
