diff --git a/client/src/main/java/io/split/storages/memory/InMemoryCacheImp.java b/client/src/main/java/io/split/storages/memory/InMemoryCacheImp.java
index 1bfbd1c4a..85acbb0eb 100644
--- a/client/src/main/java/io/split/storages/memory/InMemoryCacheImp.java
+++ b/client/src/main/java/io/split/storages/memory/InMemoryCacheImp.java
@@ -206,8 +206,10 @@ private void removeFromFlagSets(String featureFlagName, HashSet<String> sets) {
         if (sets != null) {
             for (String set : sets) {
                 HashSet<String> features = _flagSets.get(set);
-                features.remove(featureFlagName);
-                _flagSets.put(set, features);
+                if (features != null) {
+                    features.remove(featureFlagName);
+                    _flagSets.put(set, features);
+                }
             }
         }
     }
diff --git a/client/src/test/java/io/split/storages/memory/InMemoryCacheTest.java b/client/src/test/java/io/split/storages/memory/InMemoryCacheTest.java
index 7f969c266..9d347f77c 100644
--- a/client/src/test/java/io/split/storages/memory/InMemoryCacheTest.java
+++ b/client/src/test/java/io/split/storages/memory/InMemoryCacheTest.java
@@ -151,7 +151,7 @@ public void testSegmentNames() {
     }
 
     private ParsedSplit getParsedSplit(String splitName) {
-        return ParsedSplit.createParsedSplitForTests(splitName, 0, false, "default_treatment", new ArrayList<>(), "tt", 123, 2, new HashSet<>());
+        return ParsedSplit.createParsedSplitForTests(splitName, 0, false, "default_treatment", new ArrayList<>(), "tt", 123, 2, new HashSet<>(Arrays.asList("set1", "set2", "set3")));
     }
 
     @Test