Improve self-test CLI to auto-remove host trailing slash

databrickslabs · Oct 23, 2024 · e077aaf · e077aaf
1 parent c409095
commit e077aaf
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/self-testing-partner-cli/src/main/scala/com/databricks/labs/DatabricksConfigParser.scala b/self-testing-partner-cli/src/main/scala/com/databricks/labs/DatabricksConfigParser.scala
@@ -27,11 +27,8 @@ object DatabricksConfigParser {
       configurationProfileName: String,
       sysEnv: Map[String, String]
   ): (Option[String], Option[String]) = {
-    if (useEnvVariables) {
-      getConfigFromEnvVariables(sysEnv)
-    } else {
-      getDatabricksConfigFromProfile(configurationProfileName, sysEnv)
-    }
+    val (host, token) = if (useEnvVariables) getConfigFromEnvVariables(sysEnv) else getDatabricksConfigFromProfile(configurationProfileName, sysEnv)
+    return (host.map(_.stripSuffix("/")), token)
   }
 
   /** Finds the value for a field in a line from a Databricks configuration

diff --git a/self-testing-partner-cli/src/test/resources/.testcfg b/self-testing-partner-cli/src/test/resources/.testcfg
@@ -1,3 +1,7 @@
 [DEFAULT]
 host = defaulthost
 token = defaulttoken
+
+[TRAILING_SLASH]
+host = https://anotherhost.com/
+token = defaulttoken
diff --git a/self-testing-partner-cli/src/test/scala/DatabricksConfigParserSuite.scala b/self-testing-partner-cli/src/test/scala/DatabricksConfigParserSuite.scala
@@ -8,6 +8,7 @@ import org.scalatest.funsuite.AnyFunSuite
 
 class DatabricksConfigParserSuite extends AnyFunSuite {
   private val testHost = "defaulthost"
+  private val trailingSlashRemovedHost = "https://anotherhost.com"
   private val testToken = "defaulttoken"
 
   test("PROFILE - Successfully gets host/token from profile") {
@@ -46,6 +47,20 @@ class DatabricksConfigParserSuite extends AnyFunSuite {
     stream.close()
   }
 
+  test("PROFILE - Successfully remove host trailing slash from profile") {
+    val profileName = "TRAILING_SLASH"
+    val stream = new java.io.ByteArrayOutputStream()
+    val env =
+      Map(configFilePathEnvVariable -> getClass.getResource(".testcfg").getPath)
+    val (host, token) = Console.withErr(stream) {
+      getHostAndToken(useEnvVariables = false, profileName, env)
+    }
+    assert(stream.toString.isEmpty)
+    assert(host.get == trailingSlashRemovedHost)
+    assert(token.get == testToken)
+    stream.close()
+  }
+
   test(
     "ENVIRONMENT VARIABLES - Successfully gets host/token from env variables"
   ) {
@@ -87,4 +102,16 @@ class DatabricksConfigParserSuite extends AnyFunSuite {
     )
     stream.close()
   }
+
+  test("ENVIRONMENT VARIABLES - Successfully remove host trailing slash") {
+    val env = Map(hostEnvVariable -> s"$trailingSlashRemovedHost/", tokenEnvVariable -> testToken)
+    val stream = new java.io.ByteArrayOutputStream()
+    val (host, token) = Console.withOut(stream) {
+      getHostAndToken(useEnvVariables = true, "default", env)
+    }
+    assert(stream.toString.isEmpty)
+    assert(host.get == trailingSlashRemovedHost)
+    assert(token.get == testToken)
+    stream.close()
+  }
 }