From a91b9b854dbb5fc2aa353647182621e8fbdd2994 Mon Sep 17 00:00:00 2001
From: TJ Zhang <tj.zhang@improving.com>
Date: Sat, 29 Jun 2024 12:50:03 -0700
Subject: [PATCH] addressing comments

---
 glide-core/src/client/mod.rs                  | 12 +++++++
 .../test/java/glide/SharedCommandTests.java   | 33 +++++++++++++++++++
 .../test/java/glide/cluster/CommandTests.java | 17 ----------
 .../java/glide/standalone/CommandTests.java   | 17 ----------
 4 files changed, 45 insertions(+), 34 deletions(-)

diff --git a/glide-core/src/client/mod.rs b/glide-core/src/client/mod.rs
index 9961e3cf1c..e59b33e88b 100644
--- a/glide-core/src/client/mod.rs
+++ b/glide-core/src/client/mod.rs
@@ -204,6 +204,7 @@ fn get_request_timeout(cmd: &Cmd, default_timeout: Duration) -> RedisResult<Opti
             .position(b"BLOCK")
             .map(|idx| get_timeout_from_cmd_arg(cmd, idx + 1, TimeUnit::Milliseconds))
             .unwrap_or(Ok(RequestTimeoutOption::ClientConfig)),
+        b"WAIT" => get_timeout_from_cmd_arg(cmd, 1, TimeUnit::Milliseconds),
         _ => Ok(RequestTimeoutOption::ClientConfig),
     }?;
 
@@ -736,6 +737,17 @@ mod tests {
                 0.857 + BLOCKING_CMD_TIMEOUT_EXTENSION
             ))
         );
+
+        let mut cmd = Cmd::new();
+        cmd.arg("WAIT").arg("500");
+        let result = get_request_timeout(&cmd, Duration::from_millis(100));
+        assert!(result.is_ok());
+        assert_eq!(
+            result.unwrap(),
+            Some(Duration::from_secs_f64(
+                0.5 + BLOCKING_CMD_TIMEOUT_EXTENSION
+            ))
+        );
     }
 
     #[test]
diff --git a/java/integTest/src/test/java/glide/SharedCommandTests.java b/java/integTest/src/test/java/glide/SharedCommandTests.java
index fbbe76c5d8..65d732348c 100644
--- a/java/integTest/src/test/java/glide/SharedCommandTests.java
+++ b/java/integTest/src/test/java/glide/SharedCommandTests.java
@@ -7576,4 +7576,37 @@ public void geosearchstore(BaseClient client) {
                                         .get());
         assertInstanceOf(RequestException.class, requestException2.getCause());
     }
+
+    @SneakyThrows
+    @ParameterizedTest(autoCloseArguments = false)
+    @MethodSource("getClients")
+    public void waitTest(BaseClient client) {
+        // setup
+        String key = UUID.randomUUID().toString();
+        long numreplicas = 1L;
+        long timeout = 400L;
+
+        assertEquals(OK, client.set(key, "value").get());
+        assertTrue(client.wait(numreplicas, timeout).get() >= (client instanceof RedisClient ? 0 : 1));
+
+        // command should fail on a negative timeout value
+        ExecutionException executionException =
+                assertThrows(ExecutionException.class, () -> client.wait(1L, -1L).get());
+        assertInstanceOf(RequestException.class, executionException.getCause());
+    }
+
+    @SneakyThrows
+    @ParameterizedTest(autoCloseArguments = false)
+    @MethodSource("getClients")
+    public void wait_timeout_check(BaseClient client) {
+        // create new client with default request timeout (250 millis)
+        try (var testClient =
+                client instanceof RedisClient
+                        ? RedisClient.CreateClient(commonClientConfig().build()).get()
+                        : RedisClusterClient.CreateClient(commonClusterClientConfig().build()).get()) {
+
+            // ensure that commands doesn't time out even if timeout > request timeout
+            assertEquals((client instanceof RedisClient ? 0 : 1), testClient.wait(1L, 300L).get());
+        }
+    }
 }
diff --git a/java/integTest/src/test/java/glide/cluster/CommandTests.java b/java/integTest/src/test/java/glide/cluster/CommandTests.java
index b412f6abae..79106ec451 100644
--- a/java/integTest/src/test/java/glide/cluster/CommandTests.java
+++ b/java/integTest/src/test/java/glide/cluster/CommandTests.java
@@ -1846,21 +1846,4 @@ public void sort() {
                         .get());
         assertArrayEquals(key2DescendingListSubset, clusterClient.lrange(key3, 0, -1).get());
     }
-
-    @SneakyThrows
-    @Test
-    public void waitTest() {
-        // setup
-        String key = UUID.randomUUID().toString();
-        long numreplicas = 1L;
-        long timeout = 1000L;
-
-        assertEquals(OK, clusterClient.set(key, "value").get());
-        assertTrue(clusterClient.wait(numreplicas, timeout).get() >= 1);
-
-        // command should fail on a negative timeout value
-        ExecutionException executionException =
-                assertThrows(ExecutionException.class, () -> clusterClient.wait(1L, -1L).get());
-        assertInstanceOf(RequestException.class, executionException.getCause());
-    }
 }
diff --git a/java/integTest/src/test/java/glide/standalone/CommandTests.java b/java/integTest/src/test/java/glide/standalone/CommandTests.java
index 37eb331bd9..d69cb817e8 100644
--- a/java/integTest/src/test/java/glide/standalone/CommandTests.java
+++ b/java/integTest/src/test/java/glide/standalone/CommandTests.java
@@ -1032,21 +1032,4 @@ public void sort() {
                         .get());
         assertArrayEquals(namesSortedByAge, regularClient.lrange(storeKey, 0, -1).get());
     }
-
-    @SneakyThrows
-    @Test
-    public void waitTest() {
-        // setup
-        String key = UUID.randomUUID().toString();
-        long numreplicas = 1L;
-        long timeout = 1000L;
-
-        assertEquals(OK, regularClient.set(key, "value").get());
-        assertTrue(regularClient.wait(numreplicas, timeout).get() >= 0);
-
-        // command should fail on a negative timeout value
-        ExecutionException executionException =
-                assertThrows(ExecutionException.class, () -> regularClient.wait(1L, -1L).get());
-        assertInstanceOf(RequestException.class, executionException.getCause());
-    }
 }