@@ -306,13 +306,16 @@ public void testDoNotRetryCircuitBreakerException() {
306306 }
307307
308308 public void testLimitConcurrentNodes () {
309- var targetShards = List .of (
310- targetShard (shard1 , node1 ),
311- targetShard (shard2 , node2 ),
312- targetShard (shard3 , node3 ),
313- targetShard (shard4 , node4 ),
314- targetShard (shard5 , node5 )
315- );
309+ final int shards = 10 ;
310+ var targetShards = new ArrayList <DataNodeRequestSender .TargetShard >(shards );
311+ for (int i = 0 ; i < shards ; i ++) {
312+ targetShards .add (
313+ targetShard (
314+ new ShardId ("index" , "n/a" , i ),
315+ DiscoveryNodeUtils .builder ("node-" + i ).roles (Set .of (DATA_HOT_NODE_ROLE )).build ()
316+ )
317+ );
318+ }
316319
317320 var concurrency = randomIntBetween (1 , 2 );
318321 AtomicInteger maxConcurrentRequests = new AtomicInteger (0 );
@@ -335,10 +338,10 @@ public void testLimitConcurrentNodes() {
335338 listener .onResponse (new DataNodeComputeResponse (DriverCompletionInfo .EMPTY , Map .of ()));
336339 });
337340 }));
338- assertThat (sent .size (), equalTo (5 ));
341+ assertThat (sent .size (), equalTo (shards ));
339342 assertThat (maxConcurrentRequests .get (), equalTo (concurrency ));
340- assertThat (response .totalShards , equalTo (5 ));
341- assertThat (response .successfulShards , equalTo (5 ));
343+ assertThat (response .totalShards , equalTo (shards ));
344+ assertThat (response .successfulShards , equalTo (shards ));
342345 assertThat (response .failedShards , equalTo (0 ));
343346 }
344347
0 commit comments