From 394922ef55c18734f9072649c2f1591efce8400a Mon Sep 17 00:00:00 2001
From: Xiaobing <61892277+klsince@users.noreply.github.com>
Date: Wed, 13 Mar 2024 09:35:53 -0700
Subject: [PATCH] missed to delete the temp dir (#12637)

---
 .../java/org/apache/pinot/integration/tests/ClusterTest.java | 2 +-
 .../tests/PartialUpsertTableRebalanceIntegrationTest.java    | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java
index d2bfba1c5e84..77e84e1e9129 100644
--- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java
+++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java
@@ -346,7 +346,7 @@ protected void stopServer() {
     for (BaseServerStarter serverStarter : _serverStarters) {
       serverStarter.stop();
     }
-    FileUtils.deleteQuietly(new File(_baseInstanceDataDir + File.separator + "PinotServer"));
+    FileUtils.deleteQuietly(new File(_baseInstanceDataDir));
     _serverStarters = null;
     _serverGrpcPort = 0;
     _serverAdminApiPort = 0;
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
index 8afcbe917713..9d35da53bf34 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
@@ -28,6 +28,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.apache.commons.io.FileUtils;
 import org.apache.helix.model.IdealState;
 import org.apache.pinot.client.ResultSetGroup;
 import org.apache.pinot.common.exception.HttpErrorStatusException;
@@ -299,12 +300,14 @@ protected void createSchemaAndTable()
   }
 
   @AfterClass
-  public void tearDown() {
+  public void tearDown()
+      throws IOException {
     stopServer();
     stopBroker();
     stopController();
     stopKafka();
     stopZk();
+    FileUtils.deleteDirectory(_tempDir);
   }
 
   @Override