test: Add integration tests for sharded pub/sub auto-resubscription

Author: byunjuneseok

src/test/java/io/lettuce/core/pubsub/PubSubCommandIntegrationTests.java

@@ -611,4 +611,68 @@ void echoAllowedInSubscriptionState() {
         pubsub.unsubscribe(channel);
     }
 
+    @Test
+    void autoResubscribeOnShardChannelUnsubscribed() throws Exception {
+        final BlockingQueue<String> subscribedChannels = LettuceFactories.newBlockingQueue();
+        final BlockingQueue<String> unsubscribedChannels = LettuceFactories.newBlockingQueue();
+
+        RedisPubSubListener<String, String> listener = new RedisPubSubAdapter<String, String>() {
+
+            @Override
+            public void ssubscribed(String channel, long count) {
+                subscribedChannels.add(channel);
+            }
+
+            @Override
+            public void sunsubscribed(String channel, long count) {
+                unsubscribedChannels.add(channel);
+            }
+
+        };
+
+        pubsub.getStatefulConnection().addListener(listener);
+        pubsub.ssubscribe(shardChannel);
+
+        assertThat(subscribedChannels.take()).isEqualTo(shardChannel);
+
+        pubsub.sunsubscribe(shardChannel);
+
+        assertThat(unsubscribedChannels.take()).isEqualTo(shardChannel);
+        assertThat(subscribedChannels.poll(50, TimeUnit.MILLISECONDS)).isNull();
+
+        pubsub.getStatefulConnection().removeListener(listener);
+    }
+
+    @Test
+    void noAutoResubscribeOnIntentionalUnsubscribe() throws Exception {
+        final BlockingQueue<String> subscribedChannels = LettuceFactories.newBlockingQueue();
+        final BlockingQueue<String> unsubscribedChannels = LettuceFactories.newBlockingQueue();
+
+        RedisPubSubListener<String, String> listener = new RedisPubSubAdapter<String, String>() {
+
+            @Override
+            public void ssubscribed(String channel, long count) {
+                subscribedChannels.add(channel);
+            }
+
+            @Override
+            public void sunsubscribed(String channel, long count) {
+                unsubscribedChannels.add(channel);
+            }
+
+        };
+
+        pubsub.getStatefulConnection().addListener(listener);
+        pubsub.ssubscribe(shardChannel);
+
+        assertThat(subscribedChannels.take()).isEqualTo(shardChannel);
+
+        pubsub.sunsubscribe(shardChannel);
+        assertThat(unsubscribedChannels.take()).isEqualTo(shardChannel);
+
+        assertThat(subscribedChannels.poll(50, TimeUnit.MILLISECONDS)).isNull();
+
+        pubsub.getStatefulConnection().removeListener(listener);
+    }
+
 }

src/test/java/io/lettuce/core/pubsub/StatefulRedisPubSubConnectionImplUnitTests.java

@@ -17,13 +17,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.*;
 
-import io.lettuce.core.protocol.AsyncCommand;
-import io.lettuce.core.pubsub.api.async.RedisPubSubAsyncCommands;
-import io.lettuce.core.pubsub.PubSubOutput;
-import io.lettuce.core.pubsub.RedisPubSubAdapter;
-
 import java.lang.reflect.Field;
-import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -132,74 +126,27 @@ void resubscribeChannelAndPatternAndShardChanelSubscription() {
 
     @Test
     void autoResubscribeListenerIsRegistered() {
-        // Verify that the connection has the markIntentionalUnsubscribe method
-        // This confirms the auto-resubscribe functionality is available
         connection.markIntentionalUnsubscribe("test-channel");
-        // If no exception is thrown, the method exists and works
         assertTrue(true);
     }
 
     @Test
     void intentionalUnsubscribeBypassesAutoResubscribe() throws Exception {
-        // Test 1: Intentional unsubscribe should NOT trigger auto-resubscribe
-
-        // Create a mock async commands to verify ssubscribe is NOT called
-        RedisPubSubAsyncCommands<String, String> mockAsync = mock(RedisPubSubAsyncCommands.class);
-        StatefulRedisPubSubConnectionImpl<String, String> spyConnection = spy(connection);
-        when(spyConnection.async()).thenReturn(mockAsync);
-
-        // Mark the channel as intentionally unsubscribed
-        spyConnection.markIntentionalUnsubscribe("test-channel");
-
-        // Use reflection to access the private endpoint and trigger sunsubscribed event
-        PubSubEndpoint<String, String> endpoint = getEndpointViaReflection(spyConnection);
-        PubSubOutput<String, String> sunsubscribeMessage = createSunsubscribeMessage("test-channel", codec);
-        endpoint.notifyMessage(sunsubscribeMessage);
+        connection.markIntentionalUnsubscribe("test-channel");
 
-        // Wait a moment for any async processing
-        Thread.sleep(50);
+        RedisPubSubListener<String, String> autoResubscribeListener = getAutoResubscribeListener(connection);
 
-        // Verify that ssubscribe was NOT called (intentional unsubscribe bypassed auto-resubscribe)
-        verify(mockAsync, never()).ssubscribe("test-channel");
+        autoResubscribeListener.sunsubscribed("test-channel", 0);
+        verify(mockedWriter, never()).write(any(io.lettuce.core.protocol.RedisCommand.class));
     }
 
     @Test
     void unintentionalUnsubscribeTriggersAutoResubscribe() throws Exception {
-        // Test 2: Unintentional unsubscribe (from Redis) should trigger auto-resubscribe
-
-        // Create a fresh connection with a mock async
-        PubSubEndpoint<String, String> mockEndpoint = mock(PubSubEndpoint.class);
-        StatefulRedisPubSubConnectionImpl<String, String> testConnection = new StatefulRedisPubSubConnectionImpl<>(mockEndpoint,
-                mockedWriter, codec, timeout);
-
-        // Create a mock async commands to verify ssubscribe IS called
-        RedisPubSubAsyncCommands<String, String> mockAsync = mock(RedisPubSubAsyncCommands.class);
-        @SuppressWarnings("unchecked")
-        RedisFuture<Void> mockFuture = mock(RedisFuture.class);
-        when(mockAsync.ssubscribe("test-channel")).thenReturn(mockFuture);
-
-        StatefulRedisPubSubConnectionImpl<String, String> spyConnection = spy(testConnection);
-        when(spyConnection.async()).thenReturn(mockAsync);
+        RedisPubSubListener<String, String> autoResubscribeListener = getAutoResubscribeListener(connection);
 
-        // Get the auto-resubscribe listener directly and trigger it
-        RedisPubSubListener<String, String> autoResubscribeListener = getAutoResubscribeListener(spyConnection);
-
-        // Do NOT mark as intentional - simulate Redis server sunsubscribe during slot movement
         autoResubscribeListener.sunsubscribed("test-channel", 0);
 
-        // Wait a moment for async processing
-        Thread.sleep(50);
-
-        // Verify that ssubscribe WAS called (auto-resubscribe triggered)
-        verify(mockAsync, times(1)).ssubscribe("test-channel");
-    }
-
-    @SuppressWarnings("unchecked")
-    private PubSubEndpoint<String, String> getEndpointViaReflection(
-            StatefulRedisPubSubConnectionImpl<String, String> connection) throws Exception {
-        Field endpointField = StatefulRedisPubSubConnectionImpl.class.getDeclaredField("endpoint");
-        endpointField.setAccessible(true);
-        return (PubSubEndpoint<String, String>) endpointField.get(connection);
+        verify(mockedWriter, times(1)).write(any(io.lettuce.core.protocol.RedisCommand.class));
     }
 
     @SuppressWarnings("unchecked")
@@ -209,13 +156,4 @@ private RedisPubSubListener<String, String> getAutoResubscribeListener(
         listenerField.setAccessible(true);
         return (RedisPubSubListener<String, String>) listenerField.get(connection);
     }
-
-    private PubSubOutput<String, String> createSunsubscribeMessage(String channel, RedisCodec<String, String> codec) {
-        PubSubOutput<String, String> output = new PubSubOutput<>(codec);
-        output.set(ByteBuffer.wrap("sunsubscribe".getBytes()));
-        output.set(ByteBuffer.wrap(channel.getBytes()));
-        output.set(0L); // count
-        return output;
-    }
-
 }