From 85321b920d8790fae620bc2e08724caab9f8ab04 Mon Sep 17 00:00:00 2001
From: Kent Yao <yaooqinn@hotmail.com>
Date: Wed, 21 Oct 2020 17:47:55 +0800
Subject: [PATCH] KyuubiServerSuite

---
 kyuubi-main/pom.xml                           |  5 ++
 .../apache/kyuubi/server/KyuubiServer.scala   | 29 ++++++-----
 .../kyuubi/server/KyuubiServerSuite.scala     | 48 +++++++++++++++++++
 3 files changed, 70 insertions(+), 12 deletions(-)
 create mode 100644 kyuubi-main/src/test/scala/org/apache/kyuubi/server/KyuubiServerSuite.scala

diff --git a/kyuubi-main/pom.xml b/kyuubi-main/pom.xml
index 3d1582350..bbc36398f 100644
--- a/kyuubi-main/pom.xml
+++ b/kyuubi-main/pom.xml
@@ -50,6 +50,11 @@
             <type>test-jar</type>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-common</artifactId>
+        </dependency>
     </dependencies>
 
     <build>
diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
index 687db18f7..ec181f20c 100644
--- a/kyuubi-main/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
+++ b/kyuubi-main/src/main/scala/org/apache/kyuubi/server/KyuubiServer.scala
@@ -29,6 +29,22 @@ import org.apache.kyuubi.util.SignalRegister
 object KyuubiServer extends Logging {
   private val zkServer = new EmbeddedZkServer()
 
+  def startServer(conf: KyuubiConf): KyuubiServer = {
+    val zkEnsemble = conf.get(HA_ZK_QUORUM)
+    if (zkEnsemble == null || zkEnsemble.isEmpty) {
+      zkServer.initialize(conf)
+      zkServer.start()
+      sys.addShutdownHook(zkServer.stop())
+      conf.set(HA_ZK_QUORUM, zkServer.getConnectString)
+    }
+
+    val server = new KyuubiServer()
+    server.initialize(conf)
+    server.start()
+    sys.addShutdownHook(server.stop())
+    server
+  }
+
   def main(args: Array[String]): Unit = {
     info(
        """
@@ -51,18 +67,7 @@ object KyuubiServer extends Logging {
       s" ${Properties.javaVersion}")
     SignalRegister.registerLogger(logger)
     val conf = new KyuubiConf().loadFileDefaults()
-    val zkEnsemble = conf.get(HA_ZK_QUORUM)
-    if (zkEnsemble == null || zkEnsemble.isEmpty) {
-      zkServer.initialize(conf)
-      zkServer.start()
-      sys.addShutdownHook(zkServer.stop())
-      conf.set(HA_ZK_QUORUM, zkServer.getConnectString)
-    }
-
-    val server = new KyuubiServer()
-    server.initialize(conf)
-    server.start()
-    sys.addShutdownHook(server.stop())
+    startServer(conf)
   }
 }
 
diff --git a/kyuubi-main/src/test/scala/org/apache/kyuubi/server/KyuubiServerSuite.scala b/kyuubi-main/src/test/scala/org/apache/kyuubi/server/KyuubiServerSuite.scala
new file mode 100644
index 000000000..2a85e4b16
--- /dev/null
+++ b/kyuubi-main/src/test/scala/org/apache/kyuubi/server/KyuubiServerSuite.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.server
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.service.ServiceState._
+
+class KyuubiServerSuite extends KyuubiFunSuite {
+
+  test("kyuubi server basic") {
+    val server = new KyuubiServer()
+    val conf = KyuubiConf()
+    assert(server.getServices.isEmpty)
+    assert(server.getServiceState === LATENT)
+    server.initialize(conf)
+    assert(server.getServiceState === INITIALIZED)
+    val backendService = server.getServices.head.asInstanceOf[KyuubiBackendService]
+    assert(backendService.getServiceState == INITIALIZED)
+    assert(backendService.getServices.forall(_.getServiceState === INITIALIZED))
+
+    server.start()
+    assert(server.getServiceState === STARTED)
+    assert(backendService.getServiceState == STARTED)
+    assert(backendService.getServices.forall(_.getServiceState === STARTED))
+
+    server.stop()
+    assert(server.getServiceState === STOPPED)
+    assert(backendService.getServiceState == STOPPED)
+    assert(backendService.getServices.forall(_.getServiceState === STOPPED))
+  }
+
+}