Refactor provisioning plumbing to use Lettuce

service/src/test/java/org/whispersystems/textsecuregcm/push/ProvisioningManagerTest.java

@@ -0,0 +1,82 @@
+package org.whispersystems.textsecuregcm.push;
+
+import com.google.protobuf.ByteString;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.mockito.ArgumentCaptor;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+import org.whispersystems.textsecuregcm.redis.RedisSingletonExtension;
+import org.whispersystems.textsecuregcm.storage.PubSubProtos;
+import org.whispersystems.textsecuregcm.websocket.ProvisioningAddress;
+
+import java.time.Duration;
+import java.util.Random;
+import java.util.function.Consumer;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.after;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+
+class ProvisioningManagerTest {
+
+  private ProvisioningManager provisioningManager;
+
+  @RegisterExtension
+  static final RedisSingletonExtension REDIS_EXTENSION = RedisSingletonExtension.builder().build();
+
+  private static final long PUBSUB_TIMEOUT_MILLIS = 1_000;
+
+  @BeforeEach
+  void setUp() throws Exception {
+    provisioningManager = new ProvisioningManager(REDIS_EXTENSION.getRedisClient(), Duration.ofSeconds(1), new CircuitBreakerConfiguration());
+    provisioningManager.start();
+  }
+
+  @AfterEach
+  void tearDown() throws Exception {
+    provisioningManager.stop();
+  }
+
+  @Test
+  void sendProvisioningMessage() {
+    final ProvisioningAddress address = new ProvisioningAddress("address", 0);
+
+    final byte[] content = new byte[16];
+    new Random().nextBytes(content);
+
+    @SuppressWarnings("unchecked") final Consumer<PubSubProtos.PubSubMessage> subscribedConsumer = mock(Consumer.class);
+
+    provisioningManager.addListener(address, subscribedConsumer);
+    provisioningManager.sendProvisioningMessage(address, content);
+
+    final ArgumentCaptor<PubSubProtos.PubSubMessage> messageCaptor =
+        ArgumentCaptor.forClass(PubSubProtos.PubSubMessage.class);
+
+    verify(subscribedConsumer, timeout(PUBSUB_TIMEOUT_MILLIS)).accept(messageCaptor.capture());
+
+    assertEquals(PubSubProtos.PubSubMessage.Type.DELIVER, messageCaptor.getValue().getType());
+    assertEquals(ByteString.copyFrom(content), messageCaptor.getValue().getContent());
+  }
+
+  @Test
+  void removeListener() {
+    final ProvisioningAddress address = new ProvisioningAddress("address", 0);
+
+    final byte[] content = new byte[16];
+    new Random().nextBytes(content);
+
+    @SuppressWarnings("unchecked") final Consumer<PubSubProtos.PubSubMessage> subscribedConsumer = mock(Consumer.class);
+
+    provisioningManager.addListener(address, subscribedConsumer);
+    provisioningManager.removeListener(address);
+    provisioningManager.sendProvisioningMessage(address, content);
+
+    // Make sure that we give the message enough time to show up (if it was going to) before declaring victory
+    verify(subscribedConsumer, after(PUBSUB_TIMEOUT_MILLIS).never()).accept(any());
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/redis/RedisSingletonExtension.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2013-2022 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.redis;
+
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
+
+import io.lettuce.core.RedisClient;
+import io.lettuce.core.api.StatefulRedisConnection;
+import java.io.IOException;
+import java.net.ServerSocket;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import redis.embedded.RedisServer;
+
+public class RedisSingletonExtension implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
+
+  private static RedisServer redisServer;
+  private RedisClient redisClient;
+
+  public static class RedisSingletonExtensionBuilder {
+
+    private RedisSingletonExtensionBuilder() {
+    }
+
+    public RedisSingletonExtension build() {
+      return new RedisSingletonExtension();
+    }
+  }
+
+  public static RedisSingletonExtensionBuilder builder() {
+    return new RedisSingletonExtensionBuilder();
+  }
+
+  @Override
+  public void beforeAll(final ExtensionContext context) throws Exception {
+    assumeFalse(System.getProperty("os.name").equalsIgnoreCase("windows"));
+
+    redisServer = RedisServer.builder()
+        .setting("appendonly no")
+        .setting("save \"\"")
+        .setting("dir " + System.getProperty("java.io.tmpdir"))
+        .port(getAvailablePort())
+        .build();
+
+    redisServer.start();
+  }
+
+  @Override
+  public void beforeEach(final ExtensionContext context) {
+    redisClient = RedisClient.create(String.format("redis://127.0.0.1:%d", redisServer.ports().get(0)));
+
+    try (final StatefulRedisConnection<String, String> connection = redisClient.connect()) {
+      connection.sync().flushall();
+    }
+  }
+
+  @Override
+  public void afterEach(final ExtensionContext context) {
+    redisClient.shutdown();
+  }
+
+  @Override
+  public void afterAll(final ExtensionContext context) {
+    if (redisServer != null) {
+      redisServer.stop();
+    }
+  }
+
+  public RedisClient getRedisClient() {
+    return redisClient;
+  }
+
+  private static int getAvailablePort() throws IOException {
+    try (ServerSocket socket = new ServerSocket(0)) {
+      socket.setReuseAddress(false);
+      return socket.getLocalPort();
+    }
+  }
+}