1 files changed, 11 insertions, 9 deletions
diff --git a/tests/conftest.py b/tests/conftest.py
index 24783c0..ab29ee4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -151,12 +151,12 @@ def skip_ifmodversion_lt(min_version: str, module_name: str):
     raise AttributeError(f"No redis module named {module_name}")
 
 
-def skip_if_redis_enterprise(func):
+def skip_if_redis_enterprise():
     check = REDIS_INFO["enterprise"] is True
     return pytest.mark.skipif(check, reason="Redis enterprise")
 
 
-def skip_ifnot_redis_enterprise(func):
+def skip_ifnot_redis_enterprise():
     check = REDIS_INFO["enterprise"] is False
     return pytest.mark.skipif(check, reason="Not running in redis enterprise")
 
@@ -324,16 +324,18 @@ def master_host(request):
     yield parts.hostname, parts.port
 
 
-def wait_for_command(client, monitor, command):
+def wait_for_command(client, monitor, command, key=None):
     # issue a command with a key name that's local to this process.
     # if we find a command with our key before the command we're waiting
     # for, something went wrong
-    redis_version = REDIS_INFO["version"]
-    if LooseVersion(redis_version) >= LooseVersion("5.0.0"):
-        id_str = str(client.client_id())
-    else:
-        id_str = f"{random.randrange(2 ** 32):08x}"
-    key = f"__REDIS-PY-{id_str}__"
+    if key is None:
+        # generate key
+        redis_version = REDIS_INFO["version"]
+        if LooseVersion(redis_version) >= LooseVersion("5.0.0"):
+            id_str = str(client.client_id())
+        else:
+            id_str = f"{random.randrange(2 ** 32):08x}"
+        key = f"__REDIS-PY-{id_str}__"
     client.get(key)
     while True:
         monitor_response = monitor.next_command()