Add unit tests for ingestion management APIs

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/KafkaIngestionBaseIT.java

@@ -138,6 +138,11 @@ protected GetIngestionStateResponse getIngestionState(String indexName) throws E
         return client().admin().indices().getIngestionState(Requests.getIngestionStateRequest(indexName)).get();
     }
 
+    protected GetIngestionStateResponse getIngestionState(String[] indexNames, int[] shards) throws ExecutionException,
+        InterruptedException {
+        return client().admin().indices().getIngestionState(Requests.getIngestionStateRequest(indexNames, shards)).get();
+    }
+
     protected PauseIngestionResponse pauseIngestion(String indexName) throws ExecutionException, InterruptedException {
         return client().admin().indices().pauseIngestion(Requests.pauseIngestionRequest(indexName)).get();
     }

plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/RemoteStoreKafkaIT.java

@@ -23,6 +23,7 @@
 
 import java.nio.file.Path;
 import java.util.Arrays;
+import java.util.concurrent.ExecutionException;
 
 import static org.hamcrest.Matchers.is;
 
@@ -212,6 +213,27 @@ public void testPauseAndResumeIngestion() throws Exception {
         waitForSearchableDocs(4, Arrays.asList(nodeB, nodeC));
     }
 
+    public void testGetIngestionState() throws ExecutionException, InterruptedException {
+        internalCluster().startClusterManagerOnlyNode();
+        internalCluster().startDataOnlyNode();
+        internalCluster().startDataOnlyNode();
+        createIndexWithDefaultSettings(1, 1);
+        ensureGreen(indexName);
+
+        GetIngestionStateResponse ingestionState = getIngestionState(new String[] { indexName }, new int[] { 0 });
+        assertEquals(0, ingestionState.getFailedShards());
+        assertEquals(1, ingestionState.getSuccessfulShards());
+        assertEquals(1, ingestionState.getTotalShards());
+        assertEquals(1, ingestionState.getShardStates().length);
+        assertEquals(0, ingestionState.getShardStates()[0].getShardId());
+        assertEquals("POLLING", ingestionState.getShardStates()[0].getPollerState());
+        assertEquals("DROP", ingestionState.getShardStates()[0].getErrorPolicy());
+        assertFalse(ingestionState.getShardStates()[0].isPollerPaused());
+
+        GetIngestionStateResponse ingestionStateForInvalidShard = getIngestionState(new String[] { indexName }, new int[] { 1 });
+        assertEquals(0, ingestionStateForInvalidShard.getTotalShards());
+    }
+
     private void verifyRemoteStoreEnabled(String node) {
         GetSettingsResponse settingsResponse = client(node).admin().indices().prepareGetSettings(indexName).get();
         String remoteStoreEnabled = settingsResponse.getIndexToSettings().get(indexName).get("index.remote_store.enabled");

server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/IngestionUpdateStateResponse.java

@@ -21,7 +21,7 @@
 import java.util.Map;
 
 /**
- * Transport response for pausing ingestion.
+ * Transport response for ingestion state updates.
  *
  * @opensearch.experimental
  */

server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/resume/ResumeIngestionRequest.java

@@ -133,7 +133,7 @@ public void writeTo(StreamOutput out) throws IOException {
      * @opensearch.experimental
      */
     @ExperimentalApi
-    public class ResetSettings implements Writeable {
+    public static class ResetSettings implements Writeable {
         private final int shard;
         private final String mode;
         private final String value;

server/src/main/java/org/opensearch/action/admin/indices/streamingingestion/state/ShardIngestionState.java

@@ -52,8 +52,8 @@ public ShardIngestionState() {
     public ShardIngestionState(StreamInput in) throws IOException {
         index = in.readString();
         shardId = in.readInt();
-        pollerState = in.readString();
-        errorPolicy = in.readString();
+        pollerState = in.readOptionalString();
+        errorPolicy = in.readOptionalString();
         isPollerPaused = in.readBoolean();
     }
 
@@ -87,12 +87,16 @@ public boolean isPollerPaused() {
         return isPollerPaused;
     }
 
+    public String getErrorPolicy() {
+        return errorPolicy;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(index);
         out.writeInt(shardId);
-        out.writeString(pollerState);
-        out.writeString(errorPolicy);
+        out.writeOptionalString(pollerState);
+        out.writeOptionalString(errorPolicy);
         out.writeBoolean(isPollerPaused);
     }
 

server/src/main/java/org/opensearch/transport/client/Requests.java

@@ -636,7 +636,9 @@ public static GetIngestionStateRequest getIngestionStateRequest(String index) {
     /**
      * Creates a get ingestion state request given list of indices.
      */
-    public static GetIngestionStateRequest getIngestionStateRequest(String[] indices) {
-        return new GetIngestionStateRequest(indices);
+    public static GetIngestionStateRequest getIngestionStateRequest(String[] indices, int[] shards) {
+        GetIngestionStateRequest request = new GetIngestionStateRequest(indices);
+        request.setShards(shards);
+        return request;
     }
 }

server/src/test/java/org/opensearch/action/admin/indices/streamingingestion/IngestionUpdateStateResponseTests.java

@@ -0,0 +1,57 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion;
+
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+public class IngestionUpdateStateResponseTests extends OpenSearchTestCase {
+
+    public void testSerialization() throws IOException {
+        IngestionStateShardFailure[] shardFailures = new IngestionStateShardFailure[] {
+            new IngestionStateShardFailure("index1", 0, "test failure") };
+        IngestionUpdateStateResponse response = new IngestionUpdateStateResponse(true, true, shardFailures, "test error");
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            response.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                IngestionUpdateStateResponse deserializedResponse = new IngestionUpdateStateResponse(in);
+                assertTrue(deserializedResponse.isAcknowledged());
+                assertTrue(deserializedResponse.isShardsAcknowledged());
+                assertNotNull(deserializedResponse.getShardFailures());
+                assertEquals(1, deserializedResponse.getShardFailures().length);
+                assertEquals("index1", deserializedResponse.getShardFailures()[0].getIndex());
+                assertEquals(0, deserializedResponse.getShardFailures()[0].getShard());
+                assertEquals("test error", deserializedResponse.getErrorMessage());
+            }
+        }
+    }
+
+    public void testShardFailureGrouping() {
+        IngestionStateShardFailure[] shardFailures = new IngestionStateShardFailure[] {
+            new IngestionStateShardFailure("index1", 0, "failure 1"),
+            new IngestionStateShardFailure("index1", 1, "failure 2"),
+            new IngestionStateShardFailure("index2", 0, "failure 3") };
+        IngestionUpdateStateResponse response = new IngestionUpdateStateResponse(true, true, shardFailures, "test error");
+
+        Map<String, List<IngestionStateShardFailure>> groupedFailures = IngestionStateShardFailure.groupShardFailuresByIndex(shardFailures);
+        assertEquals(2, groupedFailures.size());
+        assertEquals(2, groupedFailures.get("index1").size());
+        assertEquals(1, groupedFailures.get("index2").size());
+        assertEquals(0, groupedFailures.get("index1").get(0).getShard());
+        assertEquals(1, groupedFailures.get("index1").get(1).getShard());
+        assertEquals(0, groupedFailures.get("index2").get(0).getShard());
+    }
+}

server/src/test/java/org/opensearch/action/admin/indices/streamingingestion/pause/PauseIngestionRequestTests.java

@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.pause;
+
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.action.support.IndicesOptions;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+
+public class PauseIngestionRequestTests extends OpenSearchTestCase {
+
+    public void testSerialization() throws IOException {
+        String[] indices = new String[] { "index1", "index2" };
+        PauseIngestionRequest request = new PauseIngestionRequest(indices);
+        request.indicesOptions(IndicesOptions.fromOptions(randomBoolean(), randomBoolean(), randomBoolean(), randomBoolean()));
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            request.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                PauseIngestionRequest deserializedRequest = new PauseIngestionRequest(in);
+                assertArrayEquals(request.indices(), deserializedRequest.indices());
+                assertEquals(request.indicesOptions(), deserializedRequest.indicesOptions());
+            }
+        }
+    }
+
+    public void testValidation() {
+        // Test with valid indices
+        PauseIngestionRequest request = new PauseIngestionRequest(new String[] { "index1", "index2" });
+        assertNull(request.validate());
+
+        // Test with empty indices
+        PauseIngestionRequest request2 = new PauseIngestionRequest(new String[0]);
+        ActionRequestValidationException e = request2.validate();
+        assertNotNull(e);
+    }
+}

server/src/test/java/org/opensearch/action/admin/indices/streamingingestion/resume/ResumeIngestionRequestTests.java

@@ -0,0 +1,62 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.resume;
+
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.action.support.IndicesOptions;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+
+public class ResumeIngestionRequestTests extends OpenSearchTestCase {
+
+    public void testSerialization() throws IOException {
+        String[] indices = new String[] { "index1", "index2" };
+        ResumeIngestionRequest request = new ResumeIngestionRequest(indices);
+        request.indicesOptions(IndicesOptions.fromOptions(randomBoolean(), randomBoolean(), randomBoolean(), randomBoolean()));
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            request.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                ResumeIngestionRequest deserializedRequest = new ResumeIngestionRequest(in);
+                assertArrayEquals(request.indices(), deserializedRequest.indices());
+                assertEquals(request.indicesOptions(), deserializedRequest.indicesOptions());
+            }
+        }
+    }
+
+    public void testValidation() {
+        // Test with valid indices
+        ResumeIngestionRequest request1 = new ResumeIngestionRequest(new String[] { "index1", "index2" });
+        assertNull(request1.validate());
+
+        // Test with empty indices
+        ResumeIngestionRequest request2 = new ResumeIngestionRequest(new String[0]);
+        ActionRequestValidationException e = request2.validate();
+        assertNotNull(e);
+    }
+
+    public void testResetSettingsSerialization() throws IOException {
+        ResumeIngestionRequest.ResetSettings settings = new ResumeIngestionRequest.ResetSettings(1, "mode", "value");
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            settings.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                ResumeIngestionRequest.ResetSettings deserialized = new ResumeIngestionRequest.ResetSettings(in);
+                assertEquals(settings.getShard(), deserialized.getShard());
+                assertEquals(settings.getMode(), deserialized.getMode());
+                assertEquals(settings.getValue(), deserialized.getValue());
+            }
+        }
+    }
+}

server/src/test/java/org/opensearch/action/admin/indices/streamingingestion/state/GetIngestionStateRequestTests.java

@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.state;
+
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+
+public class GetIngestionStateRequestTests extends OpenSearchTestCase {
+
+    public void testSerialization() throws IOException {
+        String[] indices = new String[] { "index1", "index2" };
+        int[] shards = new int[] { 0, 1, 2 };
+        GetIngestionStateRequest request = new GetIngestionStateRequest(indices);
+        request.setShards(shards);
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            request.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                GetIngestionStateRequest deserializedRequest = new GetIngestionStateRequest(in);
+                assertArrayEquals(request.getIndex(), deserializedRequest.getIndex());
+                assertArrayEquals(request.getShards(), deserializedRequest.getShards());
+            }
+        }
+    }
+
+    public void testValidation() {
+        // Test with valid indices
+        GetIngestionStateRequest request1 = new GetIngestionStateRequest(new String[] { "index1", "index2" });
+        assertNull(request1.validate());
+
+        // Test with null indices
+        GetIngestionStateRequest request2 = new GetIngestionStateRequest((String[]) null);
+        ActionRequestValidationException e = request2.validate();
+        assertNotNull(e);
+    }
+}

server/src/test/java/org/opensearch/action/admin/indices/streamingingestion/state/GetIngestionStateResponseTests.java

@@ -0,0 +1,39 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.admin.indices.streamingingestion.state;
+
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.Collections;
+
+public class GetIngestionStateResponseTests extends OpenSearchTestCase {
+
+    public void testSerialization() throws IOException {
+        ShardIngestionState[] shardStates = new ShardIngestionState[] {
+            new ShardIngestionState("index1", 0, "POLLING", "DROP", false),
+            new ShardIngestionState("index1", 1, "PAUSED", "BLOCK", true) };
+        GetIngestionStateResponse response = new GetIngestionStateResponse(shardStates, 2, 2, 0, Collections.emptyList());
+
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            response.writeTo(out);
+
+            try (StreamInput in = out.bytes().streamInput()) {
+                GetIngestionStateResponse deserializedResponse = new GetIngestionStateResponse(in);
+                assertEquals(response.getShardStates()[0].getShardId(), deserializedResponse.getShardStates()[0].getShardId());
+                assertEquals(response.getShardStates()[1].getShardId(), deserializedResponse.getShardStates()[1].getShardId());
+                assertEquals(response.getTotalShards(), deserializedResponse.getTotalShards());
+                assertEquals(response.getSuccessfulShards(), deserializedResponse.getSuccessfulShards());
+                assertEquals(response.getFailedShards(), deserializedResponse.getFailedShards());
+            }
+        }
+    }
+}