fix tests

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)

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)

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")
 }