- when(strategy.findShard(any())).thenReturn(SHARD);
- when(factory.getStrategy(any())).thenReturn(strategy);
- when(context.getShardStrategyFactory()).thenReturn(factory);
+ doReturn(SHARD).when(strategy).findShard(any());
+ final ShardStrategyFactory factory = mock(ShardStrategyFactory.class);
+ doReturn(strategy).when(factory).getStrategy(any());
+ doReturn(factory).when(context).getShardStrategyFactory();
+
+ final Configuration config = mock(Configuration.class);
+ doReturn(Set.of(SHARD)).when(config).getAllShardNames();
+ doReturn(config).when(context).getConfiguration();
+