diff --git a/libs/command/account_client.go b/libs/command/account_client.go
index 52ed6884a..258d23aab 100644
--- a/libs/command/account_client.go
+++ b/libs/command/account_client.go
@@ -6,21 +6,6 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 )
 
-// func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) context.Context {
-// 	if v := ctx.Value(workspaceClientKey); v != nil {
-// 		panic("command.SetWorkspaceClient called twice on the same context.")
-// 	}
-// 	return context.WithValue(ctx, workspaceClientKey, w)
-// }
-
-// func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient {
-// 	v := ctx.Value(workspaceClientKey)
-// 	if v == nil {
-// 		panic("command.WorkspaceClient called without calling command.SetWorkspaceClient first.")
-// 	}
-// 	return v.(*databricks.WorkspaceClient)
-// }
-
 func SetAccountClient(ctx context.Context, a *databricks.AccountClient) context.Context {
 	if v := ctx.Value(accountClientKey); v != nil {
 		panic("command.SetAccountClient called twice on the same context")
diff --git a/libs/command/account_client_test.go b/libs/command/account_client_test.go
new file mode 100644
index 000000000..065f6a0e8
--- /dev/null
+++ b/libs/command/account_client_test.go
@@ -0,0 +1,38 @@
+package command
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCommandAccountClient(t *testing.T) {
+	ctx := context.Background()
+	client := &databricks.AccountClient{
+		Config: &config.Config{
+			AccountID: "test-account",
+		},
+	}
+
+	// Panic if AccountClient is called before SetAccountClient.
+	assert.Panics(t, func() {
+		AccountClient(ctx)
+	})
+
+	ctx = SetAccountClient(context.Background(), client)
+
+	// Multiple calls should return a pointer to the same client.
+	a := AccountClient(ctx)
+	assert.Same(t, a, AccountClient(ctx))
+
+	// The client should have the correct configuration.
+	assert.Equal(t, "test-account", AccountClient(ctx).Config.AccountID)
+
+	// Second call should panic.
+	assert.Panics(t, func() {
+		SetAccountClient(ctx, client)
+	})
+}