Fixing test and allowing duplicate calls

Author: jonathan-buttner

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

@@ -204,6 +204,11 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                     service.defaultConfigIds(),
                     is(
                         List.of(
+                            new InferenceService.DefaultConfigId(
+                                ".elser-v2-elastic",
+                                MinimalServiceSettings.sparseEmbedding(),
+                                service
+                            ),
                             new InferenceService.DefaultConfigId(
                                 ".rainbow-sprinkles-elastic",
                                 MinimalServiceSettings.chatCompletion(),
@@ -216,7 +221,8 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
 
                 PlainActionFuture<List<Model>> listener = new PlainActionFuture<>();
                 service.defaultConfigs(listener);
-                assertThat(listener.actionGet(TIMEOUT).get(0).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));
+                assertThat(listener.actionGet(TIMEOUT).get(0).getConfigurations().getInferenceEntityId(), is(".elser-v2-elastic"));
+                assertThat(listener.actionGet(TIMEOUT).get(1).getConfigurations().getInferenceEntityId(), is(".rainbow-sprinkles-elastic"));
 
                 var getModelListener = new PlainActionFuture<UnparsedModel>();
                 // persists the default endpoints
@@ -244,12 +250,22 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
             try (var service = createElasticInferenceService()) {
                 service.waitForAuthorizationToComplete(TIMEOUT);
                 assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
-                assertTrue(service.defaultConfigIds().isEmpty());
+                assertThat(
+                    service.defaultConfigIds(),
+                    is(
+                        List.of(
+                            new InferenceService.DefaultConfigId(
+                                ".elser-v2-elastic",
+                                MinimalServiceSettings.sparseEmbedding(),
+                                service
+                            )
+                        )
+                    )
+                );
                 assertThat(service.supportedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING)));
 
                 var getModelListener = new PlainActionFuture<UnparsedModel>();
                 modelRegistry.getModel(".rainbow-sprinkles-elastic", getModelListener);
-
                 var exception = expectThrows(ResourceNotFoundException.class, () -> getModelListener.actionGet(TIMEOUT));
                 assertThat(exception.getMessage(), is("Inference endpoint not found [.rainbow-sprinkles-elastic]"));
             }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/registry/ModelRegistry.java

@@ -126,11 +126,20 @@ public boolean containsDefaultConfigId(String inferenceEntityId) {
         return defaultConfigIds.containsKey(inferenceEntityId);
     }
 
+    /**
+     * Adds the default configuration information if it does not already exist internally.
+     * @param defaultConfigId the default endpoint information
+     */
+    public synchronized void putDefaultIdIfAbsent(InferenceService.DefaultConfigId defaultConfigId) {
+        defaultConfigIds.putIfAbsent(defaultConfigId.inferenceId(), defaultConfigId);
+    }
+
     /**
      * Set the default inference ids provided by the services
-     * @param defaultConfigId The default
+     * @param defaultConfigId The default endpoint information
+     * @throws IllegalStateException if the {@link InferenceService.DefaultConfigId#inferenceId()} already exists internally
      */
-    public synchronized void addDefaultIds(InferenceService.DefaultConfigId defaultConfigId) {
+    public synchronized void addDefaultIds(InferenceService.DefaultConfigId defaultConfigId) throws IllegalStateException {
         var config = defaultConfigIds.get(defaultConfigId.inferenceId());
         if (config != null) {
             throw new IllegalStateException(

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

@@ -209,7 +209,7 @@ private synchronized void setAuthorizedContent(ElasticInferenceServiceAuthorizat
 
         configuration = new Configuration(authRef.get().taskTypesAndModels.getAuthorizedTaskTypes());
 
-        defaultConfigIds().forEach(modelRegistry::addDefaultIds);
+        defaultConfigIds().forEach(modelRegistry::putDefaultIdIfAbsent);
         handleRevokedDefaultConfigs(authorizedDefaultModelIds);
     }