From 0a166e71ae7f527c82d73e6a0d13d3ebac970d52 Mon Sep 17 00:00:00 2001
From: Kent Yao <yaooqinn@hotmail.com>
Date: Thu, 6 Dec 2018 11:02:15 +0800
Subject: [PATCH] fix ut

---
 .../yaooqinn/kyuubi/SecuredFunSuite.scala     | 22 +++++++++++++++++++
 .../session/KyuubiSessionSecuredSuite.scala   |  1 -
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/kyuubi-server/src/test/scala/yaooqinn/kyuubi/SecuredFunSuite.scala b/kyuubi-server/src/test/scala/yaooqinn/kyuubi/SecuredFunSuite.scala
index 7311c7da5..ba8937254 100644
--- a/kyuubi-server/src/test/scala/yaooqinn/kyuubi/SecuredFunSuite.scala
+++ b/kyuubi-server/src/test/scala/yaooqinn/kyuubi/SecuredFunSuite.scala
@@ -17,22 +17,44 @@
 
 package yaooqinn.kyuubi
 
+import java.io.IOException
+
+import org.apache.hadoop.minikdc.MiniKdc
 import org.apache.hadoop.security.UserGroupInformation
 import org.apache.spark.{KyuubiSparkUtil, SparkConf}
 
 trait SecuredFunSuite {
 
+  var kdc: MiniKdc = null
+  val baseDir = KyuubiSparkUtil.createTempDir(namePrefix = "kyuubi-kdc")
+  try {
+    val kdcConf = MiniKdc.createConf()
+    kdcConf.setProperty(MiniKdc.INSTANCE, "KyuubiKrbServer")
+    kdcConf.setProperty(MiniKdc.ORG_NAME, "KYUUBI")
+    kdcConf.setProperty(MiniKdc.ORG_DOMAIN, "COM")
+
+    if (kdc == null) {
+      kdc = new MiniKdc(kdcConf, baseDir)
+      kdc.start()
+    }
+  } catch {
+    case e: IOException =>
+      throw new AssertionError("unable to create temporary directory: " + e.getMessage)
+  }
+
   def tryWithSecurityEnabled(block: => Unit): Unit = {
     val conf = new SparkConf(true)
     assert(!UserGroupInformation.isSecurityEnabled)
     val authType = "spark.hadoop.hadoop.security.authentication"
     try {
       conf.set(authType, "KERBEROS")
+      System.setProperty("java.security.krb5.realm", kdc.getRealm)
       UserGroupInformation.setConfiguration(KyuubiSparkUtil.newConfiguration(conf))
       assert(UserGroupInformation.isSecurityEnabled)
       block
     } finally {
       conf.remove(authType)
+      System.clearProperty("java.security.krb5.realm")
       UserGroupInformation.setConfiguration(KyuubiSparkUtil.newConfiguration(conf))
       assert(!UserGroupInformation.isSecurityEnabled)
     }
diff --git a/kyuubi-server/src/test/scala/yaooqinn/kyuubi/session/KyuubiSessionSecuredSuite.scala b/kyuubi-server/src/test/scala/yaooqinn/kyuubi/session/KyuubiSessionSecuredSuite.scala
index 4ef1fbc0e..0ab02991f 100644
--- a/kyuubi-server/src/test/scala/yaooqinn/kyuubi/session/KyuubiSessionSecuredSuite.scala
+++ b/kyuubi-server/src/test/scala/yaooqinn/kyuubi/session/KyuubiSessionSecuredSuite.scala
@@ -46,7 +46,6 @@ class KyuubiSessionSecuredSuite extends SparkFunSuite with SecuredFunSuite {
     super.afterAll()
   }
 
-
   test("secured ugi test") {
     val be = server.beService
     val sessionMgr = be.getSessionManager