[FLINK-4073] fix stability of TaskManagerFailure test

flink-yarn-tests/src/test/java/org/apache/flink/yarn/YARNSessionCapacitySchedulerITCase.java

@@ -29,12 +29,13 @@ import org.apache.flink.test.testdata.WordCountData;
 import org.apache.flink.test.util.TestBaseUtils;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.yarn.api.protocolrecords.StopContainersRequest;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.NodeReport;
-import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.api.records.NodeState;
-import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.client.api.YarnClient;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.security.NMTokenIdentifier;
@@ -249,11 +250,16 @@ public class YARNSessionCapacitySchedulerITCase extends YarnTestBase {
 			List<NodeReport> nodeReports = yc.getNodeReports(NodeState.RUNNING);
 
 			// we asked for one node with 2 vcores so we expect 2 vcores
-			int userVcores = 0;
+			// note that the JobManager may also run on the NodeManager
+			boolean foundVCoresSetting = false;
 			for (NodeReport rep: nodeReports) {
-				userVcores += rep.getUsed().getVirtualCores();
+				Resource resource = rep.getUsed();
+				if (resource != null && resource.getVirtualCores() == 2) {
+					foundVCoresSetting = true;
+					break;
+				}
 			}
-			Assert.assertEquals(2, userVcores);
+			Assert.assertTrue(foundVCoresSetting);
 		} catch (Exception e) {
 			Assert.fail("Test failed: " + e.getMessage());
 		}