From 1b7a20619eb4a8831087eda1d44487314ec007dc Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Tue, 11 Jul 2023 23:17:35 -0400
Subject: [PATCH] Add tools for testing asynchronous Redis operations

---
 .../tests/util/MockRedisFuture.java           | 44 +++++++++++++++++++
 .../tests/util/RedisClusterHelper.java        | 22 +++++++---
 2 files changed, 61 insertions(+), 5 deletions(-)
 create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MockRedisFuture.java

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MockRedisFuture.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MockRedisFuture.java
new file mode 100644
index 000000000..fa9110cfd
--- /dev/null
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/MockRedisFuture.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.tests.util;
+
+import io.lettuce.core.RedisFuture;
+import org.jetbrains.annotations.NotNull;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+public class MockRedisFuture<T> extends CompletableFuture<T> implements RedisFuture<T> {
+
+  public static <T> MockRedisFuture<T> completedFuture(final T value) {
+    final MockRedisFuture<T> future = new MockRedisFuture<T>();
+    future.complete(value);
+    return future;
+  }
+
+  public static <U> MockRedisFuture<U> failedFuture(final Throwable cause) {
+    final MockRedisFuture<U> future = new MockRedisFuture<U>();
+    future.completeExceptionally(cause);
+    return future;
+  }
+
+  @Override
+  public String getError() {
+    return null;
+  }
+
+  @Override
+  public boolean await(final long l, final TimeUnit timeUnit) throws InterruptedException {
+    return false;
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java
index 3112b2d48..2ce6db09f 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java
@@ -27,6 +27,7 @@ public class RedisClusterHelper {
   @SuppressWarnings("unchecked")
   private static FaultTolerantRedisCluster buildMockRedisCluster(
       final RedisAdvancedClusterCommands<String, String> stringCommands,
+      final RedisAdvancedClusterAsyncCommands<String, String> stringAsyncCommands,
       final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands,
       final RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands,
       final RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands) {
@@ -35,6 +36,7 @@ public class RedisClusterHelper {
     final StatefulRedisClusterConnection<byte[], byte[]> binaryConnection = mock(StatefulRedisClusterConnection.class);
 
     when(stringConnection.sync()).thenReturn(stringCommands);
+    when(stringConnection.async()).thenReturn(stringAsyncCommands);
     when(binaryConnection.sync()).thenReturn(binaryCommands);
     when(binaryConnection.async()).thenReturn(binaryAsyncCommands);
     when(binaryConnection.reactive()).thenReturn(binaryReactiveCommands);
@@ -82,11 +84,16 @@ public class RedisClusterHelper {
   public static class Builder {
 
     private RedisAdvancedClusterCommands<String, String> stringCommands = mock(RedisAdvancedClusterCommands.class);
+    private RedisAdvancedClusterAsyncCommands<String, String> stringAsyncCommands =
+        mock(RedisAdvancedClusterAsyncCommands.class);
+
     private RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands = mock(RedisAdvancedClusterCommands.class);
-    private RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands = mock(
-        RedisAdvancedClusterAsyncCommands.class);
-    private RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands = mock(
-        RedisAdvancedClusterReactiveCommands.class);
+
+    private RedisAdvancedClusterAsyncCommands<byte[], byte[]> binaryAsyncCommands =
+        mock(RedisAdvancedClusterAsyncCommands.class);
+
+    private RedisAdvancedClusterReactiveCommands<byte[], byte[]> binaryReactiveCommands =
+        mock(RedisAdvancedClusterReactiveCommands.class);
 
     private Builder() {
 
@@ -97,6 +104,11 @@ public class RedisClusterHelper {
       return this;
     }
 
+    public Builder stringAsyncCommands(final RedisAdvancedClusterAsyncCommands<String, String> stringAsyncCommands) {
+      this.stringAsyncCommands = stringAsyncCommands;
+      return this;
+    }
+
     public Builder binaryCommands(final RedisAdvancedClusterCommands<byte[], byte[]> binaryCommands) {
       this.binaryCommands = binaryCommands;
       return this;
@@ -114,7 +126,7 @@ public class RedisClusterHelper {
     }
 
     public FaultTolerantRedisCluster build() {
-      return RedisClusterHelper.buildMockRedisCluster(stringCommands, binaryCommands, binaryAsyncCommands,
+      return RedisClusterHelper.buildMockRedisCluster(stringCommands, stringAsyncCommands, binaryCommands, binaryAsyncCommands,
           binaryReactiveCommands);
     }
   }