Waiting for termination in more tests

Author: jonathan-buttner

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java

@@ -123,7 +123,7 @@ public void testCallingStartTwice_ThrowsAssertionException() throws InterruptedE
         assertThat(exception.getMessage(), is("start() can only be called once"));
     }
 
-    public void testIsTerminated_AfterStopFromSeparateThread() {
+    public void testIsTerminated_AfterStopFromSeparateThread() throws InterruptedException {
         var waitToShutdown = new CountDownLatch(1);
         var waitToReturnFromSend = new CountDownLatch(1);
 
@@ -154,6 +154,8 @@ public void testIsTerminated_AfterStopFromSeparateThread() {
             fail(Strings.format("Executor finished before it was signaled to shutdown: %s", e));
         }
 
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isShutdown());
         assertTrue(service.isTerminated());
     }
@@ -326,12 +328,14 @@ public void onFailure(Exception e) {
         Future<?> executorTermination = submitShutdownRequest(waitToShutdown, waitToReturnFromSend, service);
 
         executorTermination.get(TIMEOUT.millis(), TimeUnit.MILLISECONDS);
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
 
         finishedOnResponse.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
     }
 
-    public void testExecute_NotifiesNonRateLimitedTasksOfShutdown() {
+    public void testExecute_NotifiesNonRateLimitedTasksOfShutdown() throws InterruptedException {
         var service = createRequestExecutorServiceWithMocks();
 
         var requestManager = RequestManagerTests.createMockWithRateLimitingDisabled(mock(RequestSender.class), "id");
@@ -348,10 +352,12 @@ public void testExecute_NotifiesNonRateLimitedTasksOfShutdown() {
             is("Failed to send request for inference id [id] because the request executor service has been shutdown")
         );
         assertTrue(thrownException.isExecutorShutdown());
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
     }
 
-    public void testExecute_NotifiesRateLimitedTasksOfShutdown() {
+    public void testExecute_NotifiesRateLimitedTasksOfShutdown() throws InterruptedException {
         var service = createRequestExecutorServiceWithMocks();
 
         var requestManager = RequestManagerTests.createMockWithRateLimitingEnabled(mock(RequestSender.class), "id");
@@ -375,6 +381,8 @@ public void testExecute_NotifiesRateLimitedTasksOfShutdown() {
             )
         );
         assertTrue(thrownException.isExecutorShutdown());
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
     }
 
@@ -569,6 +577,8 @@ public void testChangingCapacity_ToZero_SetsQueueCapacityToUnbounded() throws IO
         service.start();
 
         executorTermination.get(TIMEOUT.millis(), TimeUnit.MILLISECONDS);
+        service.awaitTermination(TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+
         assertTrue(service.isTerminated());
         assertThat(service.remainingQueueCapacity(requestManager), is(Integer.MAX_VALUE));
     }