From 660fe85147f6f40103d1d4a67cca1ac4fa802f38 Mon Sep 17 00:00:00 2001 From: Prabhu Joseph Date: Thu, 25 Apr 2019 01:28:12 +0530 Subject: [PATCH] YARN-6272 --- .../yarn/client/api/impl/TestAMRMClient.java | 34 +++++++++++++++++----- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClient.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClient.java index 24666f0..5f5c496 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClient.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClient.java @@ -19,6 +19,7 @@ package org.apache.hadoop.yarn.client.api.impl; import static org.assertj.core.api.Assertions.assertThat; +import org.assertj.core.api.Condition; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; @@ -1010,14 +1011,33 @@ private void doContainerResourceChange( // we should get decrease confirmation right away List updatedContainers = allocResponse.getUpdatedContainers(); - assertEquals(1, updatedContainers.size()); + + Condition expectedResult = new Condition() { + @Override + public boolean matches(Integer value) { + return value == 1 || value == 2; + } + }; + assertThat(updatedContainers.size()).is(expectedResult); + // we should get increase allocation after the next NM's heartbeat to RM - triggerSchedulingWithNMHeartBeat(); - // get allocations - allocResponse = amClient.allocate(0.1f); - updatedContainers = - allocResponse.getUpdatedContainers(); - assertEquals(1, updatedContainers.size()); + if (updatedContainers.size() == 1) { + int allocateAttempts = 0; + RMContext context = yarnCluster.getResourceManager().getRMContext(); + RMNode rmNode = context.getRMNodes().get(container3.getNodeId()); + while (allocateAttempts < 1000) { + context.getScheduler().handle(new NodeUpdateSchedulerEvent(rmNode)); + // get allocations + allocResponse = amClient.allocate(0.1f); + if (allocResponse.getUpdatedContainers().size() == 1) { + break; + } else { + sleep(20); + allocateAttempts++; + } + } + assertThat(allocResponse.getUpdatedContainers()).hasSize(1); + } } @Test -- 2.7.4 (Apple Git-66)