Trying to get request checks to work

Author: jonathan-buttner

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java

@@ -198,25 +198,30 @@ static List<UnparsedModel> getEisEndpoints(ModelRegistry modelRegistry) {
     }
 
     private void restartPollingTaskAndWaitForAuthResponse() throws Exception {
-        restartPollingTaskAndWaitForAuthResponse(admin(), authorizationTaskExecutor);
+        restartPollingTaskAndWaitForAuthResponse(admin(), authorizationTaskExecutor, webServer);
     }
 
-    static void restartPollingTaskAndWaitForAuthResponse(AdminClient adminClient, AuthorizationTaskExecutor authTaskExecutor)
-        throws Exception {
+    static void restartPollingTaskAndWaitForAuthResponse(
+        AdminClient adminClient,
+        AuthorizationTaskExecutor authTaskExecutor,
+        MockWebServer mockWebServer
+    ) throws Exception {
         cancelAuthorizationTask(adminClient);
 
         // wait for the new task to be recreated and an authorization response to be processed
         waitForAuthorizationToComplete(authTaskExecutor);
+
+        assertBusy(() -> {
+            var requests = mockWebServer.requests();
+            assertThat(requests.size(), is(1));
+        });
     }
 
     static void waitForAuthorizationToComplete(AuthorizationTaskExecutor authTaskExecutor) throws Exception {
         assertBusy(() -> {
             var newPoller = authTaskExecutor.getCurrentPollerTask();
             assertNotNull(newPoller);
             newPoller.waitForAuthorizationToComplete(TimeValue.THIRTY_SECONDS);
-
-            var requests = webServer.requests();
-            assertThat(requests.size(), is(1));
         });
     }
 

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/CCMServiceIT.java

@@ -35,6 +35,7 @@
 import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
 import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.AUTHORIZED_RAINBOW_SPRINKLES_RESPONSE;
 import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.AUTH_TASK_ACTION;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.EMPTY_AUTH_RESPONSE;
 import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.assertChatCompletionEndpointExists;
 import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.getEisEndpoints;
 import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.removeEisPreconfiguredEndpoints;
@@ -165,13 +166,17 @@ public void testCreatesEisChatCompletionEndpoint() throws Exception {
         waitForTask(AUTH_TASK_ACTION, admin());
         waitForAuthorizationToComplete(authorizationTaskExecutor);
 
+        assertHasRequestWithAuth();
+
+        assertChatCompletionEndpointExists(modelRegistry);
+    }
+
+    private void assertHasRequestWithAuth() throws Exception {
         assertBusy(() -> {
             var requests = webServer.requests();
             assertThat(requests.size(), is(1));
             assertThat(requests.get(0).getHeader(HttpHeaders.AUTHORIZATION), is("Bearer secret"));
         });
-
-        assertChatCompletionEndpointExists(modelRegistry);
     }
 
     private void forceClusterUpdate() {
@@ -184,6 +189,8 @@ public void testDisableCCM_RemovesAuthorizationTask() throws Exception {
         assertCCMDisabled();
         waitForNoTask(AUTH_TASK_ACTION, admin());
 
+        webServer.clearRequests();
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EMPTY_AUTH_RESPONSE));
         var listener = new TestPlainActionFuture<Void>();
         ccmService.get().storeConfiguration(new CCMModel(new SecureString("secret".toCharArray())), listener);
         listener.actionGet(TimeValue.THIRTY_SECONDS);
@@ -194,6 +201,8 @@ public void testDisableCCM_RemovesAuthorizationTask() throws Exception {
         waitForTask(AUTH_TASK_ACTION, admin());
         waitForAuthorizationToComplete(authorizationTaskExecutor);
 
+        assertHasRequestWithAuth();
+
         disableCCM();
         assertCCMDisabled();
         waitForNoTask(AUTH_TASK_ACTION, admin());