[hotfix][yarn][test] Avoid using Mockito for TaskExecutorGateway in Y…

…arnResourceManagerTest.
bbejeck · Apr 25, 2020 · 47025f7 · 47025f7
1 parent a864a8e
commit 47025f7
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java b/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerTest.java
@@ -53,8 +53,8 @@
 import org.apache.flink.runtime.rpc.TestingRpcService;
 import org.apache.flink.runtime.taskexecutor.SlotReport;
 import org.apache.flink.runtime.taskexecutor.SlotStatus;
-import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway;
 import org.apache.flink.runtime.taskexecutor.TaskExecutorRegistrationSuccess;
+import org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder;
 import org.apache.flink.runtime.util.TestingFatalErrorHandler;
 import org.apache.flink.util.ExceptionUtils;
 import org.apache.flink.util.TestLogger;
@@ -115,7 +115,6 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.mock;
 
 /**
  * General tests for the YARN resource manager component.
@@ -376,8 +375,7 @@ public void testStopWorker() throws Exception {
  verifyFutureCompleted(startContainerAsyncFuture);
 
  // Remote task executor registers with YarnResourceManager.
- TaskExecutorGateway mockTaskExecutorGateway = mock(TaskExecutorGateway.class);
- rpcService.registerGateway(taskHost, mockTaskExecutorGateway);
+ rpcService.registerGateway(taskHost, new TestingTaskExecutorGatewayBuilder().createTestingTaskExecutorGateway());
 
  final ResourceManagerGateway rmGateway = resourceManager.getSelfGateway(ResourceManagerGateway.class);