iter

Author: jimczi

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
 import org.elasticsearch.xpack.inference.external.http.HttpClientManager;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
@@ -57,7 +58,7 @@ public void createComponents() throws Exception {
         threadPool = createThreadPool(inferenceUtilityPool());
         webServer.start();
         gatewayUrl = getUrl(webServer);
-        modelRegistry = new ModelRegistry(client());
+        modelRegistry = node().injector().getInstance(ModelRegistry.class);
     }
 
     @After
@@ -73,7 +74,7 @@ protected boolean resetNodeAfterTest() {
 
     @Override
     protected Collection<Class<? extends Plugin>> getPlugins() {
-        return pluginList(ReindexPlugin.class);
+        return pluginList(ReindexPlugin.class, LocalStateInferencePlugin.class);
     }
 
     public void testDefaultConfigs_Returns_DefaultChatCompletion_V1_WhenTaskTypeIsCorrect() throws Exception {
@@ -208,7 +209,11 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                     service.defaultConfigIds(),
                     is(
                         List.of(
-                            new InferenceService.DefaultConfigId(".elser-v2-elastic", MinimalServiceSettings.sparseEmbedding(), service),
+                            new InferenceService.DefaultConfigId(
+                                ".elser-v2-elastic",
+                                MinimalServiceSettings.sparseEmbedding(ElasticInferenceService.NAME),
+                                service
+                            ),
                             new InferenceService.DefaultConfigId(
                                 ".rainbow-sprinkles-elastic",
                                 MinimalServiceSettings.chatCompletion(ElasticInferenceService.NAME),
@@ -254,7 +259,11 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
                     service.defaultConfigIds(),
                     is(
                         List.of(
-                            new InferenceService.DefaultConfigId(".elser-v2-elastic", MinimalServiceSettings.sparseEmbedding(), service)
+                            new InferenceService.DefaultConfigId(
+                                ".elser-v2-elastic",
+                                MinimalServiceSettings.sparseEmbedding(ElasticInferenceService.NAME),
+                                service
+                            )
                         )
                     )
                 );