Add unit test

Author: benchaplin

server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java

@@ -41,6 +41,8 @@
 import org.elasticsearch.cluster.node.VersionInformation;
 import org.elasticsearch.cluster.project.TestProjectResolvers;
 import org.elasticsearch.cluster.routing.IndexRoutingTable;
+import org.elasticsearch.cluster.routing.OperationRouting;
+import org.elasticsearch.cluster.routing.ShardIterator;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.TestShardRouting;
@@ -133,6 +135,8 @@
 import static org.hamcrest.Matchers.hasSize;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -1812,4 +1816,106 @@ public void onFailure(Exception ex) {
             assertTrue(ESTestCase.terminate(threadPool));
         }
     }
+
+    public void testSkippedIndicesWithRefreshBlock() {
+        final ProjectId projectId = randomProjectIdOrDefault();
+
+        String normalIndexName = "test-normal";
+        String blockedIndexName = "test-blocked";
+        final String[] indexNames = {normalIndexName, blockedIndexName};
+        final Index normalIndex = new Index(normalIndexName, UUIDs.randomBase64UUID());
+        final Index blockedIndex = new Index(blockedIndexName, UUIDs.randomBase64UUID());
+        final int numberOfShards = randomIntBetween(1, 3);
+        final int numberOfReplicas = randomIntBetween(0, 1);
+        final int totalShards = numberOfShards + numberOfShards * numberOfReplicas;
+
+        ClusterState clusterState = ClusterStateCreationUtils.stateWithAssignedPrimariesAndReplicas(
+            projectId,
+            indexNames,
+            numberOfShards,
+            numberOfReplicas
+        );
+        ClusterBlocks.Builder blocksBuilder = ClusterBlocks.builder().blocks(clusterState.blocks());
+        blocksBuilder.addIndexBlock(projectId, "test-blocked", IndexMetadata.INDEX_REFRESH_BLOCK);
+        clusterState = ClusterState.builder(clusterState).blocks(blocksBuilder).build();
+        List<ShardIterator> shardIts = new ArrayList<>();
+        for (int i = 0; i < totalShards; i++) {
+            shardIts.add(new ShardIterator(new ShardId(normalIndex, randomInt()), Collections.emptyList()));
+            shardIts.add(new ShardIterator(new ShardId(blockedIndex, randomInt()), Collections.emptyList()));
+        }
+        final OperationRouting operationRouting = mock(OperationRouting.class);
+        when(
+            operationRouting.searchShards(
+                eq(clusterState.projectState(projectId)),
+                eq(indexNames),
+                any(),
+                nullable(String.class),
+                any(),
+                any()
+            )
+        ).thenReturn(shardIts);
+        ClusterService clusterService = mock(ClusterService.class);
+        when(clusterService.state()).thenReturn(clusterState);
+        when(clusterService.getSettings()).thenReturn(Settings.EMPTY);
+        when(clusterService.operationRouting()).thenReturn(operationRouting);
+
+        Settings settings = Settings.builder()
+            .put("node.name", TransportSearchAction.class.getSimpleName())
+            .build();
+        TransportVersion transportVersion = TransportVersionUtils.getNextVersion(TransportVersions.MINIMUM_CCS_VERSION, true);
+        ThreadPool threadPool = new ThreadPool(settings, MeterRegistry.NOOP, new DefaultBuiltInExecutorBuilders());
+        try {
+            TransportService transportService = MockTransportService.createNewService(
+                Settings.EMPTY,
+                VersionInformation.CURRENT,
+                transportVersion,
+                threadPool
+            );
+            NodeClient client = new NodeClient(settings, threadPool);
+            SearchService searchService = mock(SearchService.class);
+            when(searchService.getRewriteContext(any(), any(), any(), anyBoolean())).thenReturn(
+                new QueryRewriteContext(null, null, null, null, null, null)
+            );
+
+            TransportSearchAction transportSearchAction = new TransportSearchAction(
+                threadPool,
+                new NoneCircuitBreakerService(),
+                transportService,
+                searchService,
+                null,
+                new SearchTransportService(transportService, client, null),
+                null,
+                clusterService,
+                new ActionFilters(Collections.emptySet()),
+                TestProjectResolvers.usingRequestHeader(threadPool.getThreadContext()),
+                TestIndexNameExpressionResolver.newInstance(threadPool.getThreadContext()),
+                null,
+                null,
+                new SearchResponseMetrics(TelemetryProvider.NOOP.getMeterRegistry()),
+                client,
+                new UsageService()
+            );
+
+            SearchRequest searchRequest = new SearchRequest(indexNames);
+            searchRequest.allowPartialSearchResults(true);
+            List<SearchShardIterator> searchShardIts = transportSearchAction.getLocalShardsIterator(
+                clusterState.projectState(projectId),
+                searchRequest,
+                searchRequest.getLocalClusterAlias(),
+                new HashSet<>(),
+                indexNames
+            );
+
+            assertThat(searchShardIts.size(), equalTo(shardIts.size()));
+            for (SearchShardIterator searchShardIt : searchShardIts) {
+                if (searchShardIt.skip()) {
+                    assertThat(searchShardIt.shardId().getIndexName(), equalTo("test-blocked"));
+                } else {
+                    assertThat(searchShardIt.shardId().getIndexName(), equalTo("test-normal"));
+                }
+            }
+        } finally {
+            assertTrue(ESTestCase.terminate(threadPool));
+        }
+    }
 }