From 4da4674d7ae4becd03a29354ddcc8698fc5a0e22 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 6 Mar 2025 13:02:55 +0100
Subject: [PATCH] fix tests

---
 cmd/auth/token_test.go                      |  3 ++-
 cmd/configure/configure_test.go             |  3 ++-
 cmd/root/user_agent_command_exec_id_test.go | 17 +++++------------
 3 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go
index df98cc151..917b6458d 100644
--- a/cmd/auth/token_test.go
+++ b/cmd/auth/token_test.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd"
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/cache"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
@@ -106,7 +107,7 @@ func getCobraCmdForTest(f fixtures.HTTPFixture) (*cobra.Command, *bytes.Buffer)
 func TestTokenCmdWithProfilePrintsHelpfulLoginMessageOnRefreshFailure(t *testing.T) {
 	cmd, output := getCobraCmdForTest(refreshFailureTokenResponse)
 	cmd.SetArgs([]string{"auth", "token", "--profile", "expired"})
-	err := cmd.Execute()
+	err := root.Execute(cmd.Context(), cmd)
 
 	out := output.String()
 	assert.Empty(t, out)
diff --git a/cmd/configure/configure_test.go b/cmd/configure/configure_test.go
index 14eb0674a..633944984 100644
--- a/cmd/configure/configure_test.go
+++ b/cmd/configure/configure_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/cmd"
+	"github.com/databricks/cli/cmd/root"
 	"github.com/stretchr/testify/assert"
 	"gopkg.in/ini.v1"
 )
@@ -91,7 +92,7 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
 	cmd := cmd.New(ctx)
 	cmd.SetArgs([]string{"configure", "--token", "--host", "https://host"})
 
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	_, err = os.Stat(cfgPath)
diff --git a/cmd/root/user_agent_command_exec_id_test.go b/cmd/root/user_agent_command_exec_id_test.go
index 5c4365107..a106e161f 100644
--- a/cmd/root/user_agent_command_exec_id_test.go
+++ b/cmd/root/user_agent_command_exec_id_test.go
@@ -2,25 +2,18 @@ package root
 
 import (
 	"context"
-	"regexp"
 	"testing"
 
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/databricks-sdk-go/useragent"
-	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 func TestWithCommandExecIdInUserAgent(t *testing.T) {
-	ctx := withCommandExecIdInUserAgent(context.Background())
+	ctx := command.MockExecId(context.Background(), "some-exec-id")
+	ctx = withCommandExecIdInUserAgent(ctx)
 
-	// Check that the command exec ID is in the user agent string.
+	// Check that the command exec ID is set in the user agent string.
 	ua := useragent.FromContext(ctx)
-	re := regexp.MustCompile(`cmd-exec-id/([a-f0-9-]+)`)
-	matches := re.FindAllStringSubmatch(ua, -1)
-
-	// Assert that we have exactly one match and that it's a valid UUID.
-	require.Len(t, matches, 1)
-	_, err := uuid.Parse(matches[0][1])
-	assert.NoError(t, err)
+	assert.Contains(t, ua, "cmd-exec-id/some-exec-id")
 }