From 52e24c2972f05fd7dd27b87e3c12bbb16aea3b8c Mon Sep 17 00:00:00 2001 From: Mahesh Raju Somalaraju Date: Tue, 14 Jul 2026 14:03:13 +0530 Subject: [PATCH] TEZ-4736: Fix flaky TestMRRJobsDAGApi.testMultipleMRRSleepJobViaSession --- .../apache/tez/mapreduce/TestMRRJobsDAGApi.java | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java b/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java index 92392ea6a3..94bd8d5cdb 100644 --- a/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java +++ b/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java @@ -400,16 +400,24 @@ public void testAMRelocalization() throws Exception { additionalResources.put("test.jar", createLrObjFromPath(relocFilePath)); additionalResources.put("TezAppJar.jar", createLrObjFromPath(tezAppJarRemote)); - assertEquals(TezAppMasterStatus.READY, tezSession.getAppMasterStatus()); + waitForSessionReady(tezSession); finalState = testMRRSleepJobDagSubmitCore(true, false, false, tezSession, true, MRInputAMSplitGeneratorRelocalizationTest.class, additionalResources); assertEquals(DAGStatus.State.SUCCEEDED, finalState); - assertEquals(TezAppMasterStatus.READY, tezSession.getAppMasterStatus()); + waitForSessionReady(tezSession); assertTrue(remoteFs.exists(new Path("/tmp/relocalizationfilefound"))); stopAndVerifyYarnApp(tezSession); } + private void waitForSessionReady(TezClient tezSession) + throws IOException, TezException, InterruptedException { + boolean ready = tezSession.waitTillReady(60, TimeUnit.SECONDS); + assertTrue(ready, "TezSession did not reach READY within 60s, current status: " + + tezSession.getAppMasterStatus()); + assertEquals(TezAppMasterStatus.READY, tezSession.getAppMasterStatus()); + } + private void stopAndVerifyYarnApp(TezClient tezSession) throws TezException, IOException, YarnException { ApplicationId appId = tezSession.getAppMasterApplicationId(); @@ -512,11 +520,11 @@ public void testMultipleMRRSleepJobViaSession() throws IOException, State finalState = testMRRSleepJobDagSubmitCore(true, false, false, tezSession, false, null, null); assertEquals(DAGStatus.State.SUCCEEDED, finalState); - assertEquals(TezAppMasterStatus.READY, tezSession.getAppMasterStatus()); + waitForSessionReady(tezSession); finalState = testMRRSleepJobDagSubmitCore(true, false, false, tezSession, false, null, null); assertEquals(DAGStatus.State.SUCCEEDED, finalState); - assertEquals(TezAppMasterStatus.READY, tezSession.getAppMasterStatus()); + waitForSessionReady(tezSession); stopAndVerifyYarnApp(tezSession); }