[FLINK-17636][network][tests] Fix the unstable unit tests in SingleIn…

…putGateTest
silenceland · May 14, 2020 · c781d17 · c781d17
1 parent 1d9d0bf
commit c781d17
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/...test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java b/...test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java
@@ -30,7 +30,6 @@
 import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
 import org.apache.flink.runtime.event.TaskEvent;
-import org.apache.flink.runtime.execution.CancelTaskException;
 import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
 import org.apache.flink.runtime.io.network.NettyShuffleEnvironment;
 import org.apache.flink.runtime.io.network.NettyShuffleEnvironmentBuilder;
@@ -226,7 +225,10 @@ public void testConcurrentReadStateAndProcessAndClose() throws Exception {
 						if (!getNextBufferAndVerify(inputGate, states)) {
 							Thread.sleep(1);
 						}
-					} catch (CancelTaskException expected) {
+					} catch (Throwable t) {
+						if (!inputGate.getCloseFuture().isDone()) {
+							throw new AssertionError("Exceptions are expected here only if the gate was closed", t);
+						}
 						return null;
 					}
 				}