From 5a26ef58dc9793cb5f914b70387b6310c94a9cd7 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner Date: Wed, 7 Jan 2026 15:36:24 -0500 Subject: [PATCH 1/2] Fixing flaky test for auth integrationt tests --- muted-tests.yml | 3 --- .../AuthorizationTaskExecutorIT.java | 20 ++++++++++++++++--- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/muted-tests.yml b/muted-tests.yml index a162f106237a5..c32859572b739 100644 --- a/muted-tests.yml +++ b/muted-tests.yml @@ -351,9 +351,6 @@ tests: - class: org.elasticsearch.packaging.test.DockerTests method: test072RunEsAsDifferentUserAndGroup issue: https://github.com/elastic/elasticsearch/issues/140127 -- class: org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT - method: testCreatesEisChatCompletion_DoesNotRemoveEndpointWhenNoLongerAuthorized - issue: https://github.com/elastic/elasticsearch/issues/138480 # Examples: # diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java index acc478b16901d..8136d64308a83 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java @@ -7,6 +7,8 @@ package org.elasticsearch.xpack.inference.integration; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.elasticsearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequestBuilder; import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksRequestBuilder; import org.elasticsearch.action.support.PlainActionFuture; @@ -19,6 +21,7 @@ import org.elasticsearch.reindex.ReindexPlugin; import org.elasticsearch.tasks.TaskInfo; import org.elasticsearch.test.ESSingleNodeTestCase; +import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.http.MockResponse; import org.elasticsearch.test.http.MockWebServer; import org.elasticsearch.xpack.inference.LocalStateInferencePlugin; @@ -64,6 +67,7 @@ public class AuthorizationTaskExecutorIT extends ESSingleNodeTestCase { public static final String AUTH_TASK_ACTION = AuthorizationPoller.TASK_NAME + "[c]"; + private static final Logger logger = LogManager.getLogger(AuthorizationTaskExecutorIT.class); private static final MockWebServer webServer = new MockWebServer(); private static String gatewayUrl; private static String chatCompletionResponseBody; @@ -80,8 +84,6 @@ public static void initClass() throws IOException { @Before public void createComponents() { - // Adding an empty response to ensure that the initial authorization polling request does not fail - webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EIS_EMPTY_RESPONSE)); modelRegistry = node().injector().getInstance(ModelRegistry.class); authorizationTaskExecutor = node().injector().getInstance(AuthorizationTaskExecutor.class); } @@ -95,7 +97,11 @@ static void removeEisPreconfiguredEndpoints(ModelRegistry modelRegistry) { // Delete all the eis preconfigured endpoints var listener = new PlainActionFuture(); modelRegistry.deleteModels(EIS_PRECONFIGURED_ENDPOINT_IDS, listener); - listener.actionGet(TimeValue.THIRTY_SECONDS); + try { + listener.actionGet(ESTestCase.TEST_REQUEST_TIMEOUT); + } catch (Exception e) { + logger.atWarn().withThrowable(e).log("Failed to delete eis preconfigured endpoints"); + } } @AfterClass @@ -103,6 +109,14 @@ public static void cleanUpClass() { webServer.close(); } + @Override + protected void startNode(long seed) throws Exception { + // Adding an empty response to ensure that the initial authorization polling request does not fail + // We're doing this before the node is started to ensure that the authorization task isn't created yet + webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EIS_EMPTY_RESPONSE)); + super.startNode(seed); + } + @Override protected Settings nodeSettings() { return Settings.builder() From adc4055e4514249fcbf9d0db6fb0cfe959399600 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner Date: Wed, 7 Jan 2026 15:46:07 -0500 Subject: [PATCH 2/2] Resetting the response queues --- .../test/http/MockWebServer.java | 7 +++++++ .../AuthorizationTaskExecutorIT.java | 19 ++++++++++++------- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/test/http/MockWebServer.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/test/http/MockWebServer.java index 3b118a2c6eb63..558e92b3d0128 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/test/http/MockWebServer.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/test/http/MockWebServer.java @@ -306,6 +306,13 @@ public void clearRequests() { requests.clear(); } + /** + * Removes all responses from the queue. + */ + public void clearResponses() { + responses.clear(); + } + /** * A utility method to peek into the requests and find out if #MockWebServer.takeRequests will not throw an out of bound exception * @return true if more requests are available, false otherwise diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java index 8136d64308a83..2ed8102ab70dc 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java @@ -145,7 +145,7 @@ protected boolean resetNodeAfterTest() { public void testCreatesEisChatCompletionEndpoint() throws Exception { assertNoAuthorizedEisEndpoints(); - webServer.clearRequests(); + resetWebServerQueues(); webServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionResponseBody)); restartPollingTaskAndWaitForAuthResponse(); @@ -260,14 +260,14 @@ static void cancelAuthorizationTask(AdminClient adminClient) throws Exception { public void testCreatesEisChatCompletion_DoesNotRemoveEndpointWhenNoLongerAuthorized() throws Exception { assertNoAuthorizedEisEndpoints(); - webServer.clearRequests(); + resetWebServerQueues(); webServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionResponseBody)); restartPollingTaskAndWaitForAuthResponse(); assertWebServerReceivedRequest(); assertChatCompletionEndpointExists(); - webServer.clearRequests(); + resetWebServerQueues(); // Simulate that the model is no longer authorized webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EIS_EMPTY_RESPONSE)); restartPollingTaskAndWaitForAuthResponse(); @@ -276,6 +276,11 @@ public void testCreatesEisChatCompletion_DoesNotRemoveEndpointWhenNoLongerAuthor assertChatCompletionEndpointExists(); } + private static void resetWebServerQueues() { + webServer.clearRequests(); + webServer.clearResponses(); + } + private void assertChatCompletionEndpointExists() throws Exception { assertChatCompletionEndpointExists(modelRegistry); } @@ -300,7 +305,7 @@ static void assertChatCompletionUnparsedModel(UnparsedModel rainbowSprinklesMode public void testCreatesChatCompletion_AndThenCreatesTextEmbedding() throws Exception { assertNoAuthorizedEisEndpoints(); - webServer.clearRequests(); + resetWebServerQueues(); webServer.enqueue(new MockResponse().setResponseCode(200).setBody(chatCompletionResponseBody)); restartPollingTaskAndWaitForAuthResponse(); assertWebServerReceivedRequest(); @@ -308,14 +313,14 @@ public void testCreatesChatCompletion_AndThenCreatesTextEmbedding() throws Excep assertChatCompletionEndpointExists(); // Simulate that the model is no longer authorized - webServer.clearRequests(); + resetWebServerQueues(); webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EIS_EMPTY_RESPONSE)); restartPollingTaskAndWaitForAuthResponse(); assertWebServerReceivedRequest(); assertChatCompletionEndpointExists(); - webServer.clearRequests(); + resetWebServerQueues(); // Simulate that a text embedding model is now authorized var jinaEmbedResponseBody = ElasticInferenceServiceAuthorizationResponseEntityTests.getEisJinaEmbedAuthorizationResponse(gatewayUrl) .responseJson(); @@ -341,7 +346,7 @@ public void testRestartsTaskAfterAbort() throws Exception { // Ensure the task is created and we get an initial authorization response assertNoAuthorizedEisEndpoints(); - webServer.clearRequests(); + resetWebServerQueues(); webServer.enqueue(new MockResponse().setResponseCode(200).setBody(EIS_EMPTY_RESPONSE)); // Abort the task and ensure it is restarted restartPollingTaskAndWaitForAuthResponse();