JedisPubSub and BinaryJedisPubSub PING don't support optional argument (#2254)

* JedisPubSub and BinaryJedisPubSub PING don't support optional argument

* PR review changes

* Removes unused import

Co-authored-by: M Sazzadul Hoque <[email protected]>

Author: sabbey37

src/main/java/redis/clients/jedis/BinaryJedisPubSub.java

@@ -2,6 +2,7 @@
 
 import static redis.clients.jedis.Protocol.Keyword.MESSAGE;
 import static redis.clients.jedis.Protocol.Keyword.PMESSAGE;
+import static redis.clients.jedis.Protocol.Keyword.PONG;
 import static redis.clients.jedis.Protocol.Keyword.PSUBSCRIBE;
 import static redis.clients.jedis.Protocol.Keyword.PUNSUBSCRIBE;
 import static redis.clients.jedis.Protocol.Keyword.SUBSCRIBE;
@@ -34,6 +35,9 @@ public void onPUnsubscribe(byte[] pattern, int subscribedChannels) {
   public void onPSubscribe(byte[] pattern, int subscribedChannels) {
   }
 
+  public void onPong(byte[] pattern) {
+  }
+
   public void unsubscribe() {
     client.unsubscribe();
     client.flush();
@@ -64,6 +68,16 @@ public void punsubscribe(byte[]... patterns) {
     client.flush();
   }
 
+  public void ping() {
+    client.ping();
+    client.flush();
+  }
+
+  public void ping(byte[] argument) {
+    client.ping(argument);
+    client.flush();
+  }
+
   public boolean isSubscribed() {
     return subscribedChannels > 0;
   }
@@ -115,6 +129,9 @@ private void process(Client client) {
         subscribedChannels = ((Long) reply.get(2)).intValue();
         final byte[] bpattern = (byte[]) reply.get(1);
         onPUnsubscribe(bpattern, subscribedChannels);
+      } else if (Arrays.equals(PONG.raw, resp)) {
+        final byte[] bpattern = (byte[]) reply.get(1);
+        onPong(bpattern);
       } else {
         throw new JedisException("Unknown message type: " + firstObj);
       }

src/main/java/redis/clients/jedis/JedisPubSub.java

@@ -99,6 +99,14 @@ public void ping() {
     client.flush();
   }
 
+  public void ping(String argument) {
+    if (client == null) {
+      throw new JedisConnectionException(JEDIS_SUBSCRIPTION_MESSAGE);
+    }
+    client.ping(argument);
+    client.flush();
+  }
+
   public boolean isSubscribed() {
     return subscribedChannels > 0;
   }

src/test/java/redis/clients/jedis/tests/commands/PublishSubscribeCommandsTest.java

@@ -1,12 +1,15 @@
 package redis.clients.jedis.tests.commands;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import java.io.IOException;
 import java.net.UnknownHostException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -112,6 +115,41 @@ public void onUnsubscribe(String channel, int subscribedChannels) {
     assertEquals(0L, latchUnsubscribed.getCount());
   }
 
+  @Test
+  public void pubSubChannelWithPingPongWithArgument() throws InterruptedException {
+    final CountDownLatch latchUnsubscribed = new CountDownLatch(1);
+    final CountDownLatch latchReceivedPong = new CountDownLatch(1);
+    final List<String> pongPatterns = new ArrayList<>();
+    jedis.subscribe(new JedisPubSub() {
+
+      @Override
+      public void onSubscribe(String channel, int subscribedChannels) {
+        publishOne("testchan1", "hello");
+      }
+
+      @Override
+      public void onMessage(String channel, String message) {
+        this.ping("hi!");
+      }
+
+      @Override
+      public void onPong(String pattern) {
+        pongPatterns.add(pattern);
+        latchReceivedPong.countDown();
+        unsubscribe();
+      }
+
+      @Override
+      public void onUnsubscribe(String channel, int subscribedChannels) {
+        latchUnsubscribed.countDown();
+      }
+    }, "testchan1");
+
+    assertEquals(0L, latchReceivedPong.getCount());
+    assertEquals(0L, latchUnsubscribed.getCount());
+    assertEquals(Collections.singletonList("hi!"), pongPatterns);
+  }
+
   @Test
   public void pubSubNumPat() {
     jedis.psubscribe(new JedisPubSub() {
@@ -301,6 +339,75 @@ public void onPMessage(byte[] pattern, byte[] channel, byte[] message) {
     }, SafeEncoder.encode("foo.*"), SafeEncoder.encode("bar.*"));
   }
 
+  @Test
+  public void binaryPubSubChannelWithPingPong() throws InterruptedException {
+    final CountDownLatch latchUnsubscribed = new CountDownLatch(1);
+    final CountDownLatch latchReceivedPong = new CountDownLatch(1);
+
+    jedis.subscribe(new BinaryJedisPubSub() {
+
+      @Override
+      public void onSubscribe(byte[] channel, int subscribedChannels) {
+        publishOne("testchan1", "hello");
+      }
+
+      @Override
+      public void onMessage(byte[] channel, byte[] message) {
+        this.ping();
+      }
+
+      @Override
+      public void onPong(byte[] pattern) {
+        latchReceivedPong.countDown();
+        unsubscribe();
+      }
+
+      @Override
+      public void onUnsubscribe(byte[] channel, int subscribedChannels) {
+        latchUnsubscribed.countDown();
+      }
+    }, SafeEncoder.encode("testchan1"));
+    assertEquals(0L, latchReceivedPong.getCount());
+    assertEquals(0L, latchUnsubscribed.getCount());
+  }
+
+  @Test
+  public void binaryPubSubChannelWithPingPongWithArgument() throws InterruptedException {
+    final CountDownLatch latchUnsubscribed = new CountDownLatch(1);
+    final CountDownLatch latchReceivedPong = new CountDownLatch(1);
+    final List<byte[]> pongPatterns = new ArrayList<>();
+    final byte[] pingMessage = SafeEncoder.encode("hi!");
+
+    jedis.subscribe(new BinaryJedisPubSub() {
+
+      @Override
+      public void onSubscribe(byte[] channel, int subscribedChannels) {
+        publishOne("testchan1", "hello");
+      }
+
+      @Override
+      public void onMessage(byte[] channel, byte[] message) {
+        this.ping(pingMessage);
+      }
+
+      @Override
+      public void onPong(byte[] pattern) {
+        pongPatterns.add(pattern);
+        latchReceivedPong.countDown();
+        unsubscribe();
+      }
+
+      @Override
+      public void onUnsubscribe(byte[] channel, int subscribedChannels) {
+        latchUnsubscribed.countDown();
+      }
+    }, SafeEncoder.encode("testchan1"));
+
+    assertEquals(0L, latchReceivedPong.getCount());
+    assertEquals(0L, latchUnsubscribed.getCount());
+    assertArrayEquals(pingMessage, pongPatterns.get(0));
+  }
+
   @Test
   public void binarySubscribeLazily() throws UnknownHostException, IOException,
       InterruptedException {