From 7ae89809da9314b2c12b716ff6c1d140c79f7041 Mon Sep 17 00:00:00 2001
From: ChanHo Lee <chanho0325@gmail.com>
Date: Sat, 14 Sep 2024 11:51:26 +0900
Subject: [PATCH] [ZEPPELIN-6088] Fix some flaky tests in zeppelin-integration
 (#4826)

* Fix flaky SparkParagraphIT.testCancelPyspark

* Fix flaky PersonalizeActionsIT.testGraphAction
---
 .../org/apache/zeppelin/integration/PersonalizeActionsIT.java   | 2 ++
 .../java/org/apache/zeppelin/integration/SparkParagraphIT.java  | 1 +
 2 files changed, 3 insertions(+)

diff --git a/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/PersonalizeActionsIT.java b/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/PersonalizeActionsIT.java
index 6b66197963c..bc4e0277c01 100644
--- a/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/PersonalizeActionsIT.java
+++ b/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/PersonalizeActionsIT.java
@@ -19,6 +19,7 @@
 import org.apache.zeppelin.AbstractZeppelinIT;
 import org.apache.zeppelin.MiniZeppelinServer;
 import org.apache.zeppelin.WebDriverManager;
+import org.apache.zeppelin.ZeppelinITUtils;
 import org.apache.zeppelin.test.DownloadUtils;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -239,6 +240,7 @@ void testGraphAction() throws Exception {
 
       pollingWait(By.xpath(getParagraphXPath(1) +
           "//button[contains(@uib-tooltip, 'Table')]"), MAX_BROWSER_TIMEOUT_SEC).click();
+      ZeppelinITUtils.sleep(1000, false);
       assertEquals("fa fa-table", manager.getWebDriver().findElement(By.xpath(getParagraphXPath(1)
           + "//button[contains(@class," +
           "'btn btn-default btn-sm ng-binding ng-scope active')]//i")).getAttribute("class"));
diff --git a/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/SparkParagraphIT.java b/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/SparkParagraphIT.java
index a89a4c4db06..a81ab893ddc 100644
--- a/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/SparkParagraphIT.java
+++ b/zeppelin-integration/src/test/java/org/apache/zeppelin/integration/SparkParagraphIT.java
@@ -188,6 +188,7 @@ void testCancelPyspark() throws Exception {
       ZeppelinITUtils.sleep(2000, false);
       cancelParagraph(1);
       waitForParagraph(1, "ABORT");
+      ZeppelinITUtils.sleep(1000, false);
 
       assertEquals("ABORT", getParagraphStatus(1),
           "First paragraph status is " + getParagraphStatus(1));