Fix flaky TestExchangeClient.testAddLocation

Wait for the task3 to finish after closing the client
trinodb · Jan 20, 2022 · f009a46 · f009a46
1 parent c0afeb5
commit f009a46
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/core/trino-main/src/test/java/io/trino/operator/TestExchangeClient.java b/core/trino-main/src/test/java/io/trino/operator/TestExchangeClient.java
@@ -294,14 +294,14 @@ public void testAddLocation()
         assertTaskIsNotFinished(buffer, task3);
 
         exchangeClient.close();
+        buffer.whenTaskFinished(task3).get(10, SECONDS);
+        assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task2, task3);
+        assertThat(buffer.getFailedTasks().asMap()).isEmpty();
 
         assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(0).getHttpRequestState(), "not scheduled", "httpRequestState"));
         assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(1).getHttpRequestState(), "not scheduled", "httpRequestState"));
         assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(2).getHttpRequestState(), "not scheduled", "httpRequestState"));
 
-        assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task2, task3);
-        assertThat(buffer.getFailedTasks().asMap()).isEmpty();
-
         assertTrue(exchangeClient.isFinished());
     }