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

---
 cmd/auth/token_test.go          | 8 ++++----
 cmd/configure/configure_test.go | 6 +++---
 cmd/root/bundle_test.go         | 6 +++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go
index 2089dc442..f47b41990 100644
--- a/cmd/auth/token_test.go
+++ b/cmd/auth/token_test.go
@@ -129,7 +129,7 @@ func TestTokenCmdWithHostPrintsHelpfulLoginMessageOnRefreshFailure(t *testing.T)
 func TestTokenCmdInvalidResponse(t *testing.T) {
 	cmd, output := getCobraCmdForTest(refreshFailureInvalidResponse)
 	cmd.SetArgs([]string{"auth", "token", "--profile", "active"})
-	err := cmd.Execute()
+	err := root.Execute(cmd.Context(), cmd)
 
 	out := output.String()
 	assert.Empty(t, out)
@@ -140,7 +140,7 @@ func TestTokenCmdInvalidResponse(t *testing.T) {
 func TestTokenCmdOtherErrorResponse(t *testing.T) {
 	cmd, output := getCobraCmdForTest(refreshFailureOtherError)
 	cmd.SetArgs([]string{"auth", "token", "--profile", "active"})
-	err := cmd.Execute()
+	err := root.Execute(cmd.Context(), cmd)
 
 	out := output.String()
 	assert.Empty(t, out)
@@ -151,7 +151,7 @@ func TestTokenCmdOtherErrorResponse(t *testing.T) {
 func TestTokenCmdWithProfileSuccess(t *testing.T) {
 	cmd, output := getCobraCmdForTest(refreshSuccessTokenResponse)
 	cmd.SetArgs([]string{"auth", "token", "--profile", "active"})
-	err := cmd.Execute()
+	err := root.Execute(cmd.Context(), cmd)
 
 	out := output.String()
 	validateToken(t, out)
@@ -161,7 +161,7 @@ func TestTokenCmdWithProfileSuccess(t *testing.T) {
 func TestTokenCmdWithHostSuccess(t *testing.T) {
 	cmd, output := getCobraCmdForTest(refreshSuccessTokenResponse)
 	cmd.SetArgs([]string{"auth", "token", "--host", "https://accounts.cloud.databricks.com", "--account-id", "expired"})
-	err := cmd.Execute()
+	err := root.Execute(cmd.Context(), cmd)
 
 	out := output.String()
 	validateToken(t, out)
diff --git a/cmd/configure/configure_test.go b/cmd/configure/configure_test.go
index 5681ba136..309c65363 100644
--- a/cmd/configure/configure_test.go
+++ b/cmd/configure/configure_test.go
@@ -132,7 +132,7 @@ func TestEnvVarsConfigureNoInteractive(t *testing.T) {
 	cmd := cmd.New(ctx)
 	cmd.SetArgs([]string{"configure", "--token"})
 
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	_, err = os.Stat(cfgPath)
@@ -165,7 +165,7 @@ func TestEnvVarsConfigureNoArgsNoInteractive(t *testing.T) {
 	cmd := cmd.New(ctx)
 	cmd.SetArgs([]string{"configure"})
 
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	_, err = os.Stat(cfgPath)
@@ -194,7 +194,7 @@ func TestCustomProfileConfigureNoInteractive(t *testing.T) {
 	cmd := cmd.New(ctx)
 	cmd.SetArgs([]string{"configure", "--token", "--host", "https://host", "--profile", "CUSTOM"})
 
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	_, err = os.Stat(cfgPath)
diff --git a/cmd/root/bundle_test.go b/cmd/root/bundle_test.go
index 3517b02e4..ca60d3c4c 100644
--- a/cmd/root/bundle_test.go
+++ b/cmd/root/bundle_test.go
@@ -212,7 +212,7 @@ func TestTargetFlagFull(t *testing.T) {
 	cmd.SetArgs([]string{"version", "--target", "development"})
 
 	ctx := context.Background()
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	assert.Equal(t, "development", getTarget(cmd))
@@ -224,7 +224,7 @@ func TestTargetFlagShort(t *testing.T) {
 	cmd.SetArgs([]string{"version", "-t", "production"})
 
 	ctx := context.Background()
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	assert.Equal(t, "production", getTarget(cmd))
@@ -238,7 +238,7 @@ func TestTargetEnvironmentFlag(t *testing.T) {
 	cmd.SetArgs([]string{"version", "--environment", "development"})
 
 	ctx := context.Background()
-	err := cmd.ExecuteContext(ctx)
+	err := root.Execute(ctx, cmd)
 	assert.NoError(t, err)
 
 	assert.Equal(t, "development", getTarget(cmd))