Added testRemoteClusterAction

Author: Mikep86

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

@@ -7,22 +7,45 @@
 
 package org.elasticsearch.xpack.inference.integration;
 
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.client.internal.RemoteClusterClient;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.license.LicenseSettings;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.AbstractMultiClustersTestCase;
+import org.elasticsearch.transport.RemoteClusterService;
+import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceFieldsAction;
+import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.elasticsearch.xpack.inference.FakeMlPlugin;
 import org.elasticsearch.xpack.inference.LocalStateInferencePlugin;
+import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
+import org.junit.Before;
 
+import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.xpack.inference.integration.GetInferenceFieldsIT.assertInferenceFieldsMap;
+import static org.elasticsearch.xpack.inference.integration.GetInferenceFieldsIT.assertInferenceResultsMap;
+import static org.elasticsearch.xpack.inference.integration.IntegrationTestUtils.createInferenceEndpoint;
+import static org.elasticsearch.xpack.inference.integration.IntegrationTestUtils.generateSemanticTextMapping;
 import static org.hamcrest.Matchers.containsString;
 
 public class GetInferenceFieldsCrossClusterIT extends AbstractMultiClustersTestCase {
     private static final String REMOTE_CLUSTER = "cluster_a";
+    private static final String INDEX_NAME = "test-index";
+    private static final String INFERENCE_FIELD = "test-inference-field";
+    private static final String INFERENCE_ID = "test-inference-id";
+    private static final Map<String, Object> INFERENCE_ENDPOINT_SERVICE_SETTINGS = Map.of("model", "my_model", "api_key", "my_api_key");
+
+    private boolean clustersConfigured = false;
 
     @Override
     protected List<String> remoteClusterAlias() {
@@ -41,7 +64,15 @@ protected Settings nodeSettings() {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins(String clusterAlias) {
-        return List.of(LocalStateInferencePlugin.class);
+        return List.of(LocalStateInferencePlugin.class, TestInferenceServicePlugin.class, FakeMlPlugin.class);
+    }
+
+    @Before
+    public void configureClusters() throws Exception {
+        if (clustersConfigured == false) {
+            setupTwoClusters();
+            clustersConfigured = true;
+        }
     }
 
     public void testRemoteIndex() {
@@ -52,4 +83,39 @@ public void testRemoteIndex() {
         );
         assertThat(e.getMessage(), containsString("GetInferenceFieldsAction does not support remote indices"));
     }
+
+    public void testRemoteClusterAction() {
+        RemoteClusterClient remoteClusterClient = client().getRemoteClusterClient(
+            REMOTE_CLUSTER,
+            EsExecutors.DIRECT_EXECUTOR_SERVICE,
+            RemoteClusterService.DisconnectedStrategy.RECONNECT_IF_DISCONNECTED
+        );
+
+        var request = new GetInferenceFieldsAction.Request(Set.of(INDEX_NAME), Set.of(INFERENCE_FIELD), false, false, "foo");
+        PlainActionFuture<GetInferenceFieldsAction.Response> future = new PlainActionFuture<>();
+        remoteClusterClient.execute(GetInferenceFieldsAction.REMOTE_TYPE, request, future);
+
+        var response = future.actionGet(TEST_REQUEST_TIMEOUT);
+        assertInferenceFieldsMap(
+            response.getInferenceFieldsMap(),
+            Map.of(INDEX_NAME, Set.of(new GetInferenceFieldsIT.InferenceFieldAndId(INFERENCE_FIELD, INFERENCE_ID)))
+        );
+        assertInferenceResultsMap(response.getInferenceResultsMap(), Map.of(INFERENCE_ID, TextExpansionResults.class));
+    }
+
+    private void setupTwoClusters() throws IOException {
+        setupCluster(LOCAL_CLUSTER);
+        setupCluster(REMOTE_CLUSTER);
+    }
+
+    private void setupCluster(String clusterAlias) throws IOException {
+        final Client client = client(clusterAlias);
+
+        createInferenceEndpoint(client, TaskType.SPARSE_EMBEDDING, INFERENCE_ID, INFERENCE_ENDPOINT_SERVICE_SETTINGS);
+
+        int dataNodeCount = cluster(clusterAlias).numDataNodes();
+        XContentBuilder mappings = generateSemanticTextMapping(Map.of(INFERENCE_FIELD, INFERENCE_ID));
+        Settings indexSettings = indexSettings(randomIntBetween(1, dataNodeCount), 0).build();
+        assertAcked(client.admin().indices().prepareCreate(INDEX_NAME).setSettings(indexSettings).setMapping(mappings));
+    }
 }

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

@@ -434,7 +434,7 @@ private static <T extends Exception> void assertFailedRequest(
         exceptionValidator.accept(exception);
     }
 
-    private static void assertInferenceFieldsMap(
+    static void assertInferenceFieldsMap(
         Map<String, List<InferenceFieldMetadata>> inferenceFieldsMap,
         Map<String, Set<InferenceFieldAndId>> expectedInferenceFields
     ) {
@@ -458,7 +458,7 @@ private static void assertInferenceFieldsMap(
         }
     }
 
-    private static void assertInferenceResultsMap(
+    static void assertInferenceResultsMap(
         Map<String, InferenceResults> inferenceResultsMap,
         Map<String, Class<? extends InferenceResults>> expectedInferenceResults
     ) {
@@ -490,5 +490,5 @@ private static Map<String, Class<? extends InferenceResults>> filterExpectedInfe
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
     }
 
-    private record InferenceFieldAndId(String field, String inferenceId) {}
+    record InferenceFieldAndId(String field, String inferenceId) {}
 }