Update logic of test

Author: sakthivelmanii

google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java

@@ -344,12 +344,12 @@ public void testDeadlineExceededErrorWithOneRetryWithParallelRequests()
 
     // First set of request should fail with an error
     CountDownLatch failureCountDownLatch = new CountDownLatch(3);
+    CountDownLatch executionCountLatch = new CountDownLatch(3);
+    mockSpanner.freeze();
     for (int i = 0; i < 3; i++) {
-      if (i == 0) {
-        mockSpanner.freeze();
-      }
       executor.submit(
           () -> {
+            executionCountLatch.countDown();
             try {
               try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) {
                 //noinspection StatementWithEmptyBody
@@ -361,11 +361,10 @@ public void testDeadlineExceededErrorWithOneRetryWithParallelRequests()
               failureCountDownLatch.countDown();
             }
           });
-      if (i == 2) {
-        Thread.sleep(1000);
-        mockSpanner.unfreeze();
-      }
     }
+    assertTrue(executionCountLatch.await(2, TimeUnit.SECONDS));
+    Thread.sleep(1000);
+    mockSpanner.unfreeze();
 
     System.out.println(
         "---------------------------------------------------------------------------END");
@@ -374,13 +373,14 @@ public void testDeadlineExceededErrorWithOneRetryWithParallelRequests()
 
     // Second set of requests should pass
     CountDownLatch countDownLatch = new CountDownLatch(3);
+    CountDownLatch successExecutionCountLatch = new CountDownLatch(3);
+
+    mockSpanner.freeze();
     for (int i = 0; i < 3; i++) {
-      if (i == 0) {
-        mockSpanner.freeze();
-      }
       executor.submit(
           () -> {
             try {
+              successExecutionCountLatch.countDown();
               try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) {
                 //noinspection StatementWithEmptyBody
                 while (resultSet.next()) {
@@ -391,11 +391,10 @@ public void testDeadlineExceededErrorWithOneRetryWithParallelRequests()
               countDownLatch.countDown();
             }
           });
-      if (i == 2) {
-        Thread.sleep(1000);
-        mockSpanner.unfreeze();
-      }
     }
+    assertTrue(successExecutionCountLatch.await(2, TimeUnit.SECONDS));
+    Thread.sleep(1000);
+    mockSpanner.unfreeze();
 
     countDownLatch.await(2, TimeUnit.SECONDS);
     assertEquals(3, countDownLatch.getCount());