1 files changed, 1 insertions, 2 deletions
diff --git a/core/src/test/java/org/elasticsearch/discovery/zen/NodeRemovalClusterStateTaskExecutorTests.java b/core/src/test/java/org/elasticsearch/discovery/zen/NodeRemovalClusterStateTaskExecutorTests.java
index 1e8954330c..4492cdb52e 100644
--- a/core/src/test/java/org/elasticsearch/discovery/zen/NodeRemovalClusterStateTaskExecutorTests.java
+++ b/core/src/test/java/org/elasticsearch/discovery/zen/NodeRemovalClusterStateTaskExecutorTests.java
@@ -26,7 +26,6 @@ import org.elasticsearch.cluster.ClusterStateTaskExecutor;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.allocation.AllocationService;
-import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.transport.LocalTransportAddress;
 import org.elasticsearch.test.ESTestCase;
 
@@ -133,7 +132,7 @@ public class NodeRemovalClusterStateTaskExecutorTests extends ESTestCase {
 
         final AllocationService allocationService = mock(AllocationService.class);
         when(allocationService.deassociateDeadNodes(any(ClusterState.class), eq(true), any(String.class)))
-            .thenReturn(mock(RoutingAllocation.Result.class));
+            .thenAnswer(im -> im.getArguments()[0]);
 
         final BiFunction<ClusterState, String, ClusterState> rejoin = (cs, r) -> {
             fail("rejoin should not be invoked");