From 950ea76d090e29395607ae75b49d412ba464ab63 Mon Sep 17 00:00:00 2001
From: Rong Rong <rongr@apache.org>
Date: Wed, 16 Feb 2022 14:33:41 -0800
Subject: [PATCH] Fix segment replace test (#8209)

---
 .../core/PinotHelixResourceManagerTest.java   | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerTest.java
index dac54479ec2f..1a1d080f1146 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerTest.java
@@ -96,7 +96,7 @@ public class PinotHelixResourceManagerTest {
 
   private static final int CONNECTION_TIMEOUT_IN_MILLISECOND = 10_000;
   private static final int MAXIMUM_NUMBER_OF_CONTROLLER_INSTANCES = 10;
-  private static final long TIMEOUT_IN_MS = 10_000L;
+  private static final long TIMEOUT_IN_MS = 60_000L;
 
   @BeforeClass
   public void setUp()
@@ -451,13 +451,20 @@ public void testLeadControllerAssignment() {
 
   @Test
   public void testSegmentReplacement()
-      throws IOException {
+      throws Exception {
     // Create broker tenant on 1 Brokers
     Tenant brokerTenant = new Tenant(TenantRole.BROKER, BROKER_TENANT_NAME, 1, 0, 0);
     PinotResourceManagerResponse response =
         ControllerTestUtils.getHelixResourceManager().createBrokerTenant(brokerTenant);
     Assert.assertTrue(response.isSuccessful());
 
+    testSegmentReplacementRegular();
+    testSegmentReplacementForRefresh();
+  }
+
+  private void testSegmentReplacementRegular()
+      throws IOException {
+
     // Create the table
     TableConfig tableConfig =
         new TableConfigBuilder(TableType.OFFLINE).setTableName(OFFLINE_SEGMENTS_REPLACE_TEST_TABLE_NAME)
@@ -666,14 +673,8 @@ public void testSegmentReplacement()
     Assert.assertEquals(segmentLineage.getLineageEntry(lineageEntryId4).getState(), LineageEntryState.COMPLETED);
   }
 
-  @Test
-  public void testSegmentReplacementForRefresh()
+  private void testSegmentReplacementForRefresh()
       throws IOException, InterruptedException {
-    // Create broker tenant on 1 Brokers
-    Tenant brokerTenant = new Tenant(TenantRole.BROKER, BROKER_TENANT_NAME, 1, 0, 0);
-    PinotResourceManagerResponse response =
-        ControllerTestUtils.getHelixResourceManager().createBrokerTenant(brokerTenant);
-    Assert.assertTrue(response.isSuccessful());
 
     // Create the table
     TableConfig tableConfig =