[hotfix][tests] Expose full heartbeat payload

arne-klein · Nov 29, 2019 · acae782 · acae782
1 parent 85905f8
commit acae782
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/...st/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java b/...st/java/org/apache/flink/runtime/resourcemanager/utils/TestingResourceManagerGateway.java
@@ -87,7 +87,7 @@ public class TestingResourceManagerGateway implements ResourceManagerGateway {
 
 	private volatile Function<Tuple3<ResourceID, InstanceID, SlotReport>, CompletableFuture<Acknowledge>> sendSlotReportFunction;
 
-	private volatile BiConsumer<ResourceID, SlotReport> taskExecutorHeartbeatConsumer;
+	private volatile BiConsumer<ResourceID, TaskExecutorHeartbeatPayload> taskExecutorHeartbeatConsumer;
 
 	private volatile Consumer<Tuple3<InstanceID, SlotID, AllocationID>> notifySlotAvailableConsumer;
 
@@ -153,7 +153,7 @@ public void setSendSlotReportFunction(Function<Tuple3<ResourceID, InstanceID, Sl
 		this.sendSlotReportFunction = sendSlotReportFunction;
 	}
 
-	public void setTaskExecutorHeartbeatConsumer(BiConsumer<ResourceID, SlotReport> taskExecutorHeartbeatConsumer) {
+	public void setTaskExecutorHeartbeatConsumer(BiConsumer<ResourceID, TaskExecutorHeartbeatPayload> taskExecutorHeartbeatConsumer) {
 		this.taskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;
 	}
 
@@ -248,10 +248,10 @@ public CompletableFuture<Integer> getNumberOfRegisteredTaskManagers() {
 
 	@Override
 	public void heartbeatFromTaskManager(ResourceID heartbeatOrigin, TaskExecutorHeartbeatPayload heartbeatPayload) {
-		final BiConsumer<ResourceID, SlotReport> currentTaskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;
+		final BiConsumer<ResourceID, TaskExecutorHeartbeatPayload> currentTaskExecutorHeartbeatConsumer = taskExecutorHeartbeatConsumer;
 
 		if (currentTaskExecutorHeartbeatConsumer != null) {
-			currentTaskExecutorHeartbeatConsumer.accept(heartbeatOrigin, heartbeatPayload.getSlotReport());
+			currentTaskExecutorHeartbeatConsumer.accept(heartbeatOrigin, heartbeatPayload);
 		}
 	}
 

diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java
@@ -471,7 +471,7 @@ public void testHeartbeatSlotReporting() throws Exception {
 		});
 
 		final CompletableFuture<SlotReport> heartbeatSlotReportFuture = new CompletableFuture<>();
-		rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, slotReport) -> heartbeatSlotReportFuture.complete(slotReport));
+		rmGateway.setTaskExecutorHeartbeatConsumer((resourceID, heartbeatPayload) -> heartbeatSlotReportFuture.complete(heartbeatPayload.getSlotReport()));
 
 		rpc.registerGateway(rmAddress, rmGateway);
 
@@ -1780,9 +1780,9 @@ public void testSlotReportDoesNotContainStaleInformation() throws Exception {
 		final OneShotLatch terminateSlotReportVerification = new OneShotLatch();
 		final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
 		// Assertions for this test
-		testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, slotReport) -> {
+		testingResourceManagerGateway.setTaskExecutorHeartbeatConsumer((ignored, heartbeatPayload) -> {
 			try {
-				final ArrayList<SlotStatus> slots = Lists.newArrayList(slotReport);
+				final ArrayList<SlotStatus> slots = Lists.newArrayList(heartbeatPayload.getSlotReport());
 				assertThat(slots, hasSize(1));
 				final SlotStatus slotStatus = slots.get(0);