helix-commits mailing list archives

This is an automated email from the ASF dual-hosted git repository.

hulee pushed a commit to branch restTestBase
in repository https://gitbox.apache.org/repos/asf/helix.git

commit f664bdc2fcd4594acdf7c3392e8bad0ac2c4ea72
Author: Hunter Lee <hulee@linkedin.com>
AuthorDate: Fri Jan 31 23:34:22 2020 -0800

    Upgrade ZkTestBase with multi-ZK support in helix-rest
---
 .../helix/rest/server/AbstractTestClass.java       | 38 ++++++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/helix-rest/src/test/java/org/apache/helix/rest/server/AbstractTestClass.java
b/helix-rest/src/test/java/org/apache/helix/rest/server/AbstractTestClass.java
index 11a450d..899e0b8 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/server/AbstractTestClass.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/server/AbstractTestClass.java
@@ -93,6 +93,12 @@ import org.testng.annotations.AfterSuite;
 import org.testng.annotations.BeforeSuite;
 
 public class AbstractTestClass extends JerseyTestNg.ContainerPerClassTest {
+  private static final String MULTI_ZK_PROPERTY_KEY = "multiZk";
+  private static final String NUM_ZK_PROPERTY_KEY = "numZk";
+  private static final String ZK_PREFIX = "localhost:";
+  private static final int ZK_START_PORT = 2123;
+  protected Map<String, ZkServer> _zkServerMap;
+
   protected static final String ZK_ADDR = "localhost:2123";
   protected static final String WORKFLOW_PREFIX = "Workflow_";
   protected static final String JOB_PREFIX = "Job_";
@@ -149,20 +155,43 @@ public class AbstractTestClass extends JerseyTestNg.ContainerPerClassTest
{
   @Override
   protected Application configure() {
     // start zk
+    _zkServerMap = new HashMap<>();
     try {
       if (_zkServer == null) {
         _zkServer = TestHelper.startZkServer(ZK_ADDR);
-        Assert.assertTrue(_zkServer != null);
+        Assert.assertNotNull(_zkServer);
+        _zkServerMap.put(ZK_ADDR, _zkServer);
         ZKClientPool.reset();
       }
 
       if (_zkServerTestNS == null) {
         _zkServerTestNS = TestHelper.startZkServer(_zkAddrTestNS);
-        Assert.assertTrue(_zkServerTestNS != null);
+        Assert.assertNotNull(_zkServerTestNS);
+        _zkServerMap.put(_zkAddrTestNS, _zkServerTestNS);
         ZKClientPool.reset();
       }
     } catch (Exception e) {
-      Assert.assertTrue(false, String.format("Failed to start ZK server: %s", e.toString()));
+      Assert.fail(String.format("Failed to start ZK server: %s", e.toString()));
+    }
+
+    // Start additional ZKs in a multi-ZK setup
+    String multiZkConfig = System.getProperty(MULTI_ZK_PROPERTY_KEY);
+    if (multiZkConfig != null && multiZkConfig.equalsIgnoreCase(Boolean.TRUE.toString()))
{
+      String numZkFromConfig = System.getProperty(NUM_ZK_PROPERTY_KEY);
+      if (numZkFromConfig != null) {
+        try {
+          int numZkFromConfigInt = Integer.parseInt(numZkFromConfig);
+          // Start (numZkFromConfigInt - 2) ZooKeepers
+          for (int i = 2; i < numZkFromConfigInt; i++) {
+            String zkAddr = ZK_PREFIX + (ZK_START_PORT + i);
+            ZkServer zkServer = TestHelper.startZkServer(zkAddr);
+            Assert.assertNotNull(zkServer);
+            _zkServerMap.put(zkAddr, zkServer);
+          }
+        } catch (Exception e) {
+          Assert.fail("Failed to create multiple ZooKeepers!");
+        }
+      }
     }
 
     // Configure server context
@@ -286,6 +315,9 @@ public class AbstractTestClass extends JerseyTestNg.ContainerPerClassTest
{
       _zkServerTestNS = null;
     }
 
+    // Stop all ZkServers
+    _zkServerMap.forEach((zkAddr, zkServer) -> TestHelper.stopZkServer(zkServer));
+
     if (_helixRestServer != null) {
       _helixRestServer.shutdown();
       _helixRestServer = null;