diff --git a/test/functional/api/cas/cache.py b/test/functional/api/cas/cache.py
index f9b27e2..7bd092d 100644
--- a/test/functional/api/cas/cache.py
+++ b/test/functional/api/cas/cache.py
@@ -112,7 +112,7 @@ class Cache:
     def reset_counters(self):
         return casadm.reset_counters(self.cache_id)
 
-    def set_cache_mode(self, cache_mode: CacheMode, flush: bool = True):
+    def set_cache_mode(self, cache_mode: CacheMode, flush=None):
         return casadm.set_cache_mode(cache_mode, self.cache_id, flush)
 
     def load_io_class(self, file_path: str):
diff --git a/test/functional/api/cas/casadm.py b/test/functional/api/cas/casadm.py
index 720cc0a..6f0bfa5 100644
--- a/test/functional/api/cas/casadm.py
+++ b/test/functional/api/cas/casadm.py
@@ -167,10 +167,12 @@ def print_statistics(cache_id: int, core_id: int = None, per_io_class: bool = Fa
 
 
 def set_cache_mode(cache_mode: CacheMode, cache_id: int,
-                   flush: bool = True, shortcut: bool = False):
+                   flush=None, shortcut: bool = False):
     flush_cache = None
-    if cache_mode in [CacheMode.WB, CacheMode.WO]:
-        flush_cache = "yes" if flush else "no"
+    if flush is True:
+        flush_cache = "yes"
+    elif flush is False:
+        flush_cache = "no"
 
     output = TestRun.executor.run(
         set_cache_mode_cmd(cache_mode=cache_mode.name.lower(), cache_id=str(cache_id),