From e25ed2fdbf801b814cdaa1c9bc1afcb9e79ca13d Mon Sep 17 00:00:00 2001
From: Jacky Hu <jacky.hu@databricks.com>
Date: Thu, 31 Oct 2024 09:42:53 -0700
Subject: [PATCH] Improve self-test CLI to auto-remove host trailing slash
 (#77)

---
 .../labs/DatabricksConfigParser.scala         |  7 ++---
 .../src/test/resources/.testcfg               |  4 +++
 .../scala/DatabricksConfigParserSuite.scala   | 27 +++++++++++++++++++
 3 files changed, 33 insertions(+), 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
index 01ef955..c27a131 100644
--- 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
index 279f31d..e3e2d25 100644
--- 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
index e279688..c2ff488 100644
--- 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()
+  }
 }