Speeding up test

Author: jonathan-buttner

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java

@@ -127,10 +127,10 @@ private TraceContext getCurrentTraceInfo() {
     void waitForAuthRequestCompletion(TimeValue timeValue) throws IllegalStateException {
         try {
             if (requestCompleteLatch.await(timeValue.getMillis(), TimeUnit.MILLISECONDS) == false) {
-                throw new IllegalStateException("The authorization request did not complete.");
+                throw new IllegalStateException("The wait time has expired for authorization to complete.");
             }
-        } catch (IllegalStateException | InterruptedException e) {
-            logger.warn("Interrupted while waiting for the authorization request to complete", e);
+        } catch (InterruptedException e) {
+            throw new IllegalStateException("Waiting for authorization to complete was interrupted");
         }
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandlerTests.java

@@ -194,14 +194,13 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
     }
 
     @SuppressWarnings("unchecked")
-    public void testGetAuthorization_OnResponseCalledOnce() throws IllegalStateException {
+    public void testGetAuthorization_OnResponseCalledOnce() throws IOException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var eisGatewayUrl = getUrl(webServer);
         var logger = mock(Logger.class);
         var authHandler = new ElasticInferenceServiceAuthorizationHandler(eisGatewayUrl, threadPool, logger);
 
         ActionListener<ElasticInferenceServiceAuthorization> listener = mock(ActionListener.class);
-
         String responseJson = """
                 {
                     "models": [
@@ -214,17 +213,18 @@ public void testGetAuthorization_OnResponseCalledOnce() throws IllegalStateExcep
             """;
         webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
-        authHandler.getAuthorization(listener, senderFactory.createSender());
-        authHandler.waitForAuthRequestCompletion(TimeValue.timeValueSeconds(1));
-
-        verify(listener, times(1)).onResponse(any());
+        try (var sender = senderFactory.createSender()) {
+            authHandler.getAuthorization(listener, sender);
+            authHandler.waitForAuthRequestCompletion(TIMEOUT);
 
-        var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
-        verify(logger, times(1)).debug(loggerArgsCaptor.capture());
+            verify(listener, times(1)).onResponse(any());
+            var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
+            verify(logger, times(1)).debug(loggerArgsCaptor.capture());
 
-        var message = loggerArgsCaptor.getValue();
-        assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
-        verifyNoMoreInteractions(logger);
+            var message = loggerArgsCaptor.getValue();
+            assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
+            verifyNoMoreInteractions(logger);
+        }
     }
 
     public void testGetAuthorization_InvalidResponse() throws IOException {
@@ -262,5 +262,4 @@ public void testGetAuthorization_InvalidResponse() throws IOException {
         }
 
     }
-
 }