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,17 +17,15 @@
 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;
 import java.util.List;
+import java.util.Set;
 
 @Tag(UNIT_TEST)
 class StatefulRedisPubSubConnectionImplUnitTests {
@@ -132,66 +130,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
+        RedisPubSubListener<String, String> autoResubscribeListener = getAutoResubscribeListener(connection);
 
-        // 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);
-
-        // 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");
+        verify(mockedWriter, times(1)).write(any(io.lettuce.core.protocol.RedisCommand.class));
     }
 
     @SuppressWarnings("unchecked")