Refactor and Test of ConfigSecurity

jlopezmalla · Sep 13, 2017 · b29cce4 · b29cce4
1 parent b1cad29
commit b29cce4
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 10 deletions.
diff --git a/core/src/main/scala/org/apache/spark/security/ConfigSecurity.scala b/core/src/main/scala/org/apache/spark/security/ConfigSecurity.scala
@@ -24,17 +24,22 @@ object ConfigSecurity extends Logging{
 
   var vaultToken: Option[String] = None
   val vaultHost: Option[String] = sys.env.get("VAULT_HOST")
-  val vaultUri: Option[String] = {
-    (sys.env.get("VAULT_PROTOCOL"), vaultHost, sys.env.get("VAULT_PORT")) match {
-      case (Some(vaultProtocol), Some(vaultHost), Some(vaultPort)) =>
-        val vaultUri = s"$vaultProtocol://$vaultHost:$vaultPort"
-        logDebug(s"vault uri: $vaultUri found, any Vault Connection will use it")
-        Option(vaultUri)
-      case _ =>
-        logDebug("No Vault information found, any Vault Connection will fail")
-        None
+  val vaultUri: Option[String] = getVaultUri(sys.env.get("VAULT_PROTOCOL"),
+    vaultHost, sys.env.get("VAULT_PORT"))
+
+    def getVaultUri(vaultProtocol: Option[String],
+                    vaultHost: Option[String],
+                    vaultPort: Option[String]): Option[String] = {
+      (vaultProtocol, vaultHost, vaultPort) match {
+        case (Some (vaultProtocol), Some (vaultHost), Some (vaultPort) ) =>
+          val vaultUri = s"$vaultProtocol://$vaultHost:$vaultPort"
+          logDebug (s"vault uri: $vaultUri found, any Vault Connection will use it")
+          Option (vaultUri)
+        case _ =>
+          logDebug ("No Vault information found, any Vault Connection will fail")
+          None
+      }
     }
-  }
 
   def prepareEnvironment(vaultAppToken: Option[String] = None,
                          vaulHost: Option[String] = None): Map[String, String] = {

diff --git a/core/src/test/scala/org/apache/spark/security/ConfigSecuritySuite.scala b/core/src/test/scala/org/apache/spark/security/ConfigSecuritySuite.scala
@@ -45,4 +45,28 @@ class ConfigSecuritySuite extends SparkFunSuite with Matchers {
         "datastore" -> Map("DATASTORE_ENABLE" -> "true"),
         "kafka" -> Map("KAFKA_ENABLE" -> "true")))
   }
+  /**
+    * getVaultUri
+    */
+
+  val vaultProtocol = Option("https")
+  val vaultHost = Option("vault.labs.stratio.com")
+  val vaultPort = Option("8200")
+  val expectedResult = "https://vault.labs.stratio.com:8200"
+
+  test("getVaultUri with all the parameters given") {
+      val vaultUri = ConfigSecurity.getVaultUri(vaultProtocol, vaultHost, vaultPort)
+      assert(vaultUri === Option(expectedResult))
+    }
+
+  test("getVaultUri without one of the parameters needed") {
+    val vaultUri = ConfigSecurity.getVaultUri(None, vaultHost, vaultPort)
+    assert(vaultUri === None)
+  }
+
+  test("getVaultUri without none of the parameters needed") {
+    val vaultUri = ConfigSecurity.getVaultUri(None, None, None)
+    assert(vaultUri === None)
+  }
+
 }