Add UT for listenable transport request handler

Signed-off-by: zane-neo <zaniu@amazon.com>
opensearch-project · Aug 9, 2024 · f17a5c0 · f17a5c0
1 parent 11c214f
commit f17a5c0
Showing 1 changed file with 64 additions and 0 deletions.
diff --git a/server/src/test/java/org/opensearch/action/support/nodes/TransportNodesActionTests.java b/server/src/test/java/org/opensearch/action/support/nodes/TransportNodesActionTests.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.action.support.nodes;
 
+import org.mockito.ArgumentCaptor;
 import org.opensearch.Version;
 import org.opensearch.action.FailedNodeException;
 import org.opensearch.action.support.ActionFilters;
@@ -43,16 +44,19 @@
 import org.opensearch.cluster.node.DiscoveryNodeRole;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.node.NodeService;
+import org.opensearch.tasks.Task;
 import org.opensearch.telemetry.tracing.noop.NoopTracer;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.transport.CapturingTransport;
 import org.opensearch.threadpool.TestThreadPool;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportChannel;
 import org.opensearch.transport.TransportRequest;
 import org.opensearch.transport.TransportService;
 import org.junit.After;
@@ -72,6 +76,7 @@
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.function.Supplier;
 
+import static org.mockito.Mockito.verify;
 import static org.opensearch.test.ClusterServiceUtils.createClusterService;
 import static org.opensearch.test.ClusterServiceUtils.setState;
 import static org.mockito.Mockito.mock;
@@ -166,6 +171,22 @@ public void testCustomResolving() throws Exception {
         assertEquals(clusterService.state().nodes().getDataNodes().size(), capturedRequests.size());
     }
 
+    public void testCreateTransportNodesActionWithListenableHandler() {
+        TransportNodesAction action = getListenableHandlerTestTransportNodesAction();
+        assertTrue(transport.getRequestHandlers().getHandler(action.actionName + "[n]").getHandler() instanceof TransportNodesAction.ListenableNodeTransportHandler);
+    }
+
+    public void testMessageReceivedInListenableNodeTransportHandler() throws Exception {
+        TransportNodesAction action = getListenableHandlerTestTransportNodesAction();
+        TransportChannel transportChannel = mock(TransportChannel.class);
+        transport.getRequestHandlers().getHandler(action.actionName + "[n]").getHandler().messageReceived(new TestNodeRequest(), transportChannel, mock(
+            Task.class));
+        ArgumentCaptor<TestNodeResponse> argCaptor = ArgumentCaptor.forClass(TestNodeResponse.class);
+        verify(transportChannel).sendResponse(argCaptor.capture());
+        TestNodeResponse response = argCaptor.getValue();
+        assertNotNull(response);
+    }
+
     private <T> List<T> mockList(Supplier<T> supplier, int size) {
         List<T> failures = new ArrayList<>(size);
         for (int i = 0; i < size; ++i) {
@@ -258,6 +279,19 @@ public TestTransportNodesAction getTestTransportNodesAction() {
         );
     }
 
+    public TestTransportNodesAction getListenableHandlerTestTransportNodesAction() {
+        return new TestTransportNodesAction(
+            THREAD_POOL,
+            clusterService,
+            transportService,
+            new ActionFilters(Collections.emptySet()),
+            TestNodesRequest::new,
+            TestNodeRequest::new,
+            ThreadPool.Names.SAME,
+            true
+        );
+    }
+
     public DataNodesOnlyTransportNodesAction getDataNodesOnlyTransportNodesAction(TransportService transportService) {
         return new DataNodesOnlyTransportNodesAction(
             THREAD_POOL,
@@ -303,6 +337,31 @@ private static class TestTransportNodesAction extends TransportNodesAction<
             );
         }
 
+        TestTransportNodesAction(
+            ThreadPool threadPool,
+            ClusterService clusterService,
+            TransportService transportService,
+            ActionFilters actionFilters,
+            Writeable.Reader<TestNodesRequest> request,
+            Writeable.Reader<TestNodeRequest> nodeRequest,
+            String nodeExecutor,
+            boolean listenableHandler
+        ) {
+            super(
+                "indices:admin/test",
+                threadPool,
+                clusterService,
+                transportService,
+                actionFilters,
+                request,
+                nodeRequest,
+                nodeExecutor,
+                nodeExecutor,
+                listenableHandler,
+                TestNodeResponse.class
+            );
+        }
+
         @Override
         protected TestNodesResponse newResponse(
             TestNodesRequest request,
@@ -327,6 +386,11 @@ protected TestNodeResponse nodeOperation(TestNodeRequest request) {
             return new TestNodeResponse();
         }
 
+        @Override
+        protected void nodeOperation(TestNodeRequest request, ActionListener<TestNodeResponse> actionListener) {
+            actionListener.onResponse(new TestNodeResponse());
+        }
+
     }
 
     private static class DataNodesOnlyTransportNodesAction extends TestTransportNodesAction {