acc: Add IsServicePrincipal option (#3315)

## Why
In some unit tests (e.g. run_as) we set current user to service
principal and that's important setting for those tests. This allows
converting those tests to acceptance tests.

Author: denik

acceptance/internal/config.go

@@ -52,6 +52,9 @@ type TestConfig struct {
 	// If true and Cloud=true, run this test only if a default warehouse is available in the cloud environment
 	RequiresWarehouse *bool
 
+	// If set, current user will be set to a service principal-like UUID instead of email (default is false)
+	IsServicePrincipal *bool
+
 	// List of additional replacements to apply on this test.
 	// Old is a regexp, New is a replacement expression.
 	Repls []testdiff.Replacement

acceptance/internal/handlers.go

@@ -8,7 +8,6 @@ import (
 	"strconv"
 
 	"github.com/databricks/databricks-sdk-go/service/catalog"
-	"github.com/databricks/databricks-sdk-go/service/iam"
 
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -18,11 +17,6 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-var TestUser = iam.User{
-	Id:       "1000012345",
-	UserName: "[email protected]",
-}
-
 var TestMetastore = catalog.MetastoreAssignment{
 	DefaultCatalogName: "hive_metastore",
 	MetastoreId:        "120efa64-9b68-46ba-be38-f319458430d2",
@@ -74,7 +68,7 @@ func addDefaultHandlers(server *testserver.Server) {
 	server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any {
 		return testserver.Response{
 			Headers: map[string][]string{"X-Databricks-Org-Id": {"900800700600"}},
-			Body:    TestUser,
+			Body:    req.Workspace.CurrentUser(),
 		}
 	})
 

acceptance/internal/prepare_server.go

@@ -47,7 +47,16 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o
 	// Use a unique token for each test. This allows us to maintain
 	// separate state for each test in fake workspaces.
 	tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
-	token := "dbapi" + tokenSuffix
+
+	var token string
+	var testUser iam.User
+	if isTruePtr(config.IsServicePrincipal) {
+		token = testserver.ServicePrincipalTokenPrefix + tokenSuffix
+		testUser = testserver.TestUserSP
+	} else {
+		token = testserver.UserNameTokenPrefix + tokenSuffix
+		testUser = testserver.TestUser
+	}
 
 	if cloudEnv != "" {
 		w, err := databricks.NewWorkspaceClient()
@@ -79,7 +88,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o
 			Token: token,
 		}
 
-		return cfg, TestUser
+		return cfg, testUser
 	}
 
 	// Default case. Start a dedicated local server for the test with the server stubs configured
@@ -92,7 +101,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o
 
 	// For the purposes of replacements, use testUser for local runs.
 	// Note, users might have overriden /api/2.0/preview/scim/v2/Me but that should not affect the replacement:
-	return cfg, TestUser
+	return cfg, testUser
 }
 
 func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir string) func(request *testserver.Request) {

acceptance/selftest/IsServicePrincipal/out.test.toml

@@ -0,0 +1,5 @@
+Local = true
+Cloud = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]

acceptance/selftest/IsServicePrincipal/output.txt

@@ -0,0 +1,6 @@
+
+>>> [CLI] current-user me
+{
+  "id":"[USERID]",
+  "userName":"X649126d-dc78-4e80-9c8d-86e6b3d2ffb3"
+}

acceptance/selftest/IsServicePrincipal/script

@@ -0,0 +1 @@
+trace $CLI current-user me | sed 's/e649126d/X649126d/'

acceptance/selftest/IsServicePrincipal/test.toml

@@ -0,0 +1 @@
+IsServicePrincipal = true

libs/testserver/fake_workspace.go

@@ -14,6 +14,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/apps"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/dashboards"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/databricks/databricks-sdk-go/service/sql"
@@ -26,19 +27,33 @@ import (
 // (encoding/json without options parses numbers into float64)
 // These are also easier to spot / replace in test output compared to numbers with one or few digits.
 const (
-	TestJobID = 4611686018427387911
-	TestRunID = 2305843009213693969
+	TestJobID                   = 4611686018427387911
+	TestRunID                   = 2305843009213693969
+	UserNameTokenPrefix         = "dbapi0"
+	ServicePrincipalTokenPrefix = "dbapi1"
+	UserID                      = "1000012345"
 )
 
+var TestUser = iam.User{
+	Id:       UserID,
+	UserName: "[email protected]",
+}
+
+var TestUserSP = iam.User{
+	Id:       UserID,
+	UserName: "e649126d-dc78-4e80-9c8d-86e6b3d2ffb3",
+}
+
 type FileEntry struct {
 	Info workspace.ObjectInfo
 	Data []byte
 }
 
 // FakeWorkspace holds a state of a workspace for acceptance tests.
 type FakeWorkspace struct {
-	mu  sync.Mutex
-	url string
+	mu                 sync.Mutex
+	url                string
+	isServicePrincipal bool
 
 	directories  map[string]bool
 	files        map[string]FileEntry
@@ -119,9 +134,10 @@ func MapDelete[K comparable, V any](w *FakeWorkspace, collection map[K]V, key K)
 	return Response{}
 }
 
-func NewFakeWorkspace(url string) *FakeWorkspace {
+func NewFakeWorkspace(url, token string) *FakeWorkspace {
 	return &FakeWorkspace{
-		url: url,
+		url:                url,
+		isServicePrincipal: strings.HasPrefix(token, ServicePrincipalTokenPrefix),
 		directories: map[string]bool{
 			"/Workspace": true,
 		},
@@ -144,6 +160,14 @@ func NewFakeWorkspace(url string) *FakeWorkspace {
 	}
 }
 
+func (s *FakeWorkspace) CurrentUser() iam.User {
+	if s != nil && s.isServicePrincipal {
+		return TestUserSP
+	} else {
+		return TestUser
+	}
+}
+
 func (s *FakeWorkspace) WorkspaceGetStatus(path string) Response {
 	defer s.LockUnlock()()
 

libs/testserver/pipelines_test.go

@@ -26,15 +26,15 @@ func createTestPipeline(t *testing.T, workspace *FakeWorkspace) string {
 }
 
 func TestPipelineStartUpdate_HandlesNonExistentPipeline(t *testing.T) {
-	workspace := NewFakeWorkspace("http://test")
+	workspace := NewFakeWorkspace("http://test", "dbapi123")
 
 	response := workspace.PipelineStartUpdate("non-existent-pipeline")
 	assert.Equal(t, 404, response.StatusCode)
 	assert.Contains(t, response.Body.(map[string]string)["message"], "The specified pipeline non-existent-pipeline was not found")
 }
 
 func TestPipelineGetUpdate_HandlesNonExistent(t *testing.T) {
-	workspace := NewFakeWorkspace("http://test")
+	workspace := NewFakeWorkspace("http://test", "dbapi123")
 
 	response := workspace.PipelineGetUpdate("non-existent-pipeline", "some-update-id")
 	assert.Equal(t, 404, response.StatusCode)
@@ -47,7 +47,7 @@ func TestPipelineGetUpdate_HandlesNonExistent(t *testing.T) {
 }
 
 func TestPipelineStop_AfterUpdate(t *testing.T) {
-	workspace := NewFakeWorkspace("http://test")
+	workspace := NewFakeWorkspace("http://test", "dbapi123")
 
 	pipelineId := createTestPipeline(t, workspace)
 

libs/testserver/server.go

@@ -242,7 +242,7 @@ func (s *Server) getWorkspaceForToken(token string) *FakeWorkspace {
 	defer s.mu.Unlock()
 
 	if _, ok := s.fakeWorkspaces[token]; !ok {
-		s.fakeWorkspaces[token] = NewFakeWorkspace(s.Server.URL)
+		s.fakeWorkspaces[token] = NewFakeWorkspace(s.Server.URL, token)
 	}
 
 	return s.fakeWorkspaces[token]