FixOfConnectionStateListenerTest (Azure#17481)

* fix flakyness of connectionStateListener test

Co-authored-by: Annie Liang <[email protected]>

Author: xinlian12

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ConnectionStateListenerTest.java

@@ -24,16 +24,18 @@
 import org.testng.annotations.Test;
 
 import java.util.HashMap;
+import java.util.Random;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 
 import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext;
 
 public class ConnectionStateListenerTest {
+    private static final Logger logger = LoggerFactory.getLogger(ConnectionStateListenerTest.class);
 
     private static int port = 8082;
-    private static String serverUriString = "rntbd://localhost:" + port;
-    private static final Logger logger = LoggerFactory.getLogger(ConnectionStateListenerTest.class);
+    private static String serverAddressPrefix = "rntbd://localhost:";
+    private static Random random = new Random();
 
     @DataProvider(name = "connectionStateListenerConfigProvider")
     public Object[][] connectionStateListenerConfigProvider() {
@@ -52,7 +54,9 @@ public void connectionStateListener_OnConnectionEvent(
         RequestResponseType responseType,
         int times) throws ExecutionException, InterruptedException {
 
-        TcpServer server = TcpServerFactory.startNewRntbdServer(port);
+        // using a random generated server port
+        int serverPort = port + random.nextInt(1000);
+        TcpServer server = TcpServerFactory.startNewRntbdServer(serverPort);
         // Inject fake response
         server.injectServerResponse(responseType);
 
@@ -78,17 +82,17 @@ public void connectionStateListener_OnConnectionEvent(
                 getDocumentDefinition(), new HashMap<>());
         req.setPartitionKeyRangeIdentity(new PartitionKeyRangeIdentity("fakeCollectionId","fakePartitionKeyRangeId"));
 
-        Uri targetUri = new Uri(serverUriString);
+        Uri targetUri = new Uri(serverAddressPrefix + serverPort);
         try {
             client.invokeStoreAsync(targetUri, req).block();
         } catch (Exception e) {
             logger.info("expected failed request with reason {}", e);
         }
         finally {
-            Mockito.verify(addressResolver, Mockito.times(times)).updateAddresses(Mockito.any(), Mockito.any());
+            TcpServerFactory.shutdownRntbdServer(server);
         }
 
-        TcpServerFactory.shutdownRntbdServer(server);
+        Mockito.verify(addressResolver, Mockito.times(times)).updateAddresses(Mockito.any(), Mockito.any());
     }
 
     private Document getDocumentDefinition() {

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/TcpServerMock/TcpServer.java

@@ -121,6 +121,8 @@ public void initChannel(SocketChannel channel) throws Exception {
                 logger.info("Server channel closed.");
             });
 
+        } catch (Exception e) {
+            promise.setFailure(e);
         } finally {
             parent.shutdownGracefully().sync();
             child.shutdownGracefully().sync();