Add support for mocking servers in acceptance tests

acceptance/acceptance_test.go

@@ -16,6 +16,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/databricks/cli/internal/testserver"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/testdiff"
@@ -109,11 +110,11 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	cloudEnv := os.Getenv("CLOUD_ENV")
 
 	if cloudEnv == "" {
-		server := StartServer(t)
+		defaultServer := StartServer(t)
-		AddHandlers(server)
+		AddHandlers(defaultServer)
 		// Redirect API access to local server:
-		t.Setenv("DATABRICKS_HOST", server.URL)
+		t.Setenv("DATABRICKS_HOST", defaultServer.URL)
-		t.Setenv("DATABRICKS_TOKEN", "dapi1234")
+		t.Setenv("DATABRICKS_TOKEN", "acceptance-test-token")
 
 		homeDir := t.TempDir()
 		// Do not read user's ~/.databrickscfg
@@ -143,7 +144,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	for _, dir := range testDirs {
 		testName := strings.ReplaceAll(dir, "\\", "/")
 		t.Run(testName, func(t *testing.T) {
-			if !InprocessMode {
+			if !InprocessMode && !hasCustomServer(t, dir) {
 				t.Parallel()
 			}
 
@@ -154,6 +155,10 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	return len(testDirs)
 }
 
+func hasCustomServer(t *testing.T, dir string) bool {
+	return testutil.DetectFile(t, filepath.Join(dir, "server.json"))
+}
+
 func getTests(t *testing.T) []string {
 	testDirs := make([]string, 0, 128)
 
@@ -187,6 +192,16 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 		tmpDir = t.TempDir()
 	}
 
+	// If there is a server.json file in the test directory, start a custom server.
+	// Redirect all API calls to this server.
+	if hasCustomServer(t, dir) {
+		server := testserver.NewFromConfig(t, filepath.Join(dir, "server.json"))
+		t.Setenv("DATABRICKS_HOST", server.URL)
+		t.Cleanup(func() {
+			server.Close()
+		})
+	}
+
 	// Converts C:\Users\DENIS~1.BIL -> C:\Users\denis.bilenko
 	tmpDirEvalled, err1 := filepath.EvalSymlinks(tmpDir)
 	if err1 == nil && tmpDirEvalled != tmpDir {

acceptance/bundle/templates/dbt-sql/output/my_dbt_sql/databricks.yml

@@ -3,7 +3,7 @@
 # See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
 bundle:
   name: my_dbt_sql
-  uuid: <UUID>
+  uuid: [UUID]
 
 include:
   - resources/*.yml

acceptance/bundle/templates/default-python/output/my_default_python/databricks.yml

@@ -2,7 +2,7 @@
 # See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
 bundle:
   name: my_default_python
-  uuid: <UUID>
+  uuid: [UUID]
 
 include:
   - resources/*.yml

acceptance/bundle/templates/default-python/output/my_default_python/scratch/exploration.ipynb

@@ -20,7 +20,7 @@
       "rowLimit": 10000
      },
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }

acceptance/bundle/templates/default-python/output/my_default_python/src/dlt_pipeline.ipynb

@@ -6,7 +6,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }
@@ -24,7 +24,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }
@@ -47,7 +47,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }

acceptance/bundle/templates/default-python/output/my_default_python/src/notebook.ipynb

@@ -6,7 +6,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }
@@ -37,7 +37,7 @@
       "rowLimit": 10000
      },
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }

acceptance/bundle/templates/default-sql/output/my_default_sql/databricks.yml

@@ -2,7 +2,7 @@
 # See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
 bundle:
   name: my_default_sql
-  uuid: <UUID>
+  uuid: [UUID]
 
 include:
   - resources/*.yml

acceptance/bundle/templates/default-sql/output/my_default_sql/scratch/exploration.ipynb

@@ -7,7 +7,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }

acceptance/bundle/templates/experimental-jobs-as-code/output/my_jobs_as_code/databricks.yml

@@ -2,7 +2,7 @@
 # See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
 bundle:
   name: my_jobs_as_code
-  uuid: <UUID>
+  uuid: [UUID]
 
 experimental:
   python:

acceptance/bundle/templates/experimental-jobs-as-code/output/my_jobs_as_code/src/notebook.ipynb

@@ -6,7 +6,7 @@
     "application/vnd.databricks.v1+cell": {
      "cellMetadata": {},
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }
@@ -37,7 +37,7 @@
       "rowLimit": 10000
      },
      "inputWidgets": {},
-     "nuid": "<UUID>",
+     "nuid": "[UUID]",
      "showTitle": false,
      "title": ""
     }

acceptance/cmd_server_test.go

@@ -8,10 +8,11 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/internal/testcli"
+	"github.com/databricks/cli/internal/testserver"
 	"github.com/stretchr/testify/require"
 )
 
-func StartCmdServer(t *testing.T) *TestServer {
+func StartCmdServer(t *testing.T) *testserver.Server {
 	server := StartServer(t)
 	server.Handle("/", func(r *http.Request) (any, error) {
 		q := r.URL.Query()

acceptance/server_test.go

@@ -1,74 +1,27 @@
 package acceptance_test
 
 import (
-	"encoding/json"
+	"context"
 	"net/http"
-	"net/http/httptest"
 	"testing"
 
+	"github.com/databricks/cli/internal/testserver"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-type TestServer struct {
+func StartServer(t *testing.T) *testserver.Server {
-	*httptest.Server
+	server := testserver.New(context.Background(), t)
-	Mux *http.ServeMux
-}
-
-type HandlerFunc func(r *http.Request) (any, error)
-
-func NewTestServer() *TestServer {
-	mux := http.NewServeMux()
-	server := httptest.NewServer(mux)
-
-	return &TestServer{
-		Server: server,
-		Mux:    mux,
-	}
-}
-
-func (s *TestServer) Handle(pattern string, handler HandlerFunc) {
-	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
-		resp, err := handler(r)
-		if err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
-
-		w.Header().Set("Content-Type", "application/json")
-
-		var respBytes []byte
-
-		respString, ok := resp.(string)
-		if ok {
-			respBytes = []byte(respString)
-		} else {
-			respBytes, err = json.MarshalIndent(resp, "", "    ")
-			if err != nil {
-				http.Error(w, err.Error(), http.StatusInternalServerError)
-				return
-			}
-		}
-
-		if _, err := w.Write(respBytes); err != nil {
-			http.Error(w, err.Error(), http.StatusInternalServerError)
-			return
-		}
-	})
-}
-
-func StartServer(t *testing.T) *TestServer {
-	server := NewTestServer()
 	t.Cleanup(func() {
 		server.Close()
 	})
 	return server
 }
 
-func AddHandlers(server *TestServer) {
+func AddHandlers(server *testserver.Server) {
-	server.Handle("/api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{
 				{
@@ -83,7 +36,7 @@ func AddHandlers(server *TestServer) {
 		}, nil
 	})
 
-	server.Handle("/api/2.0/instance-pools/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/instance-pools/list", func(r *http.Request) (any, error) {
 		return compute.ListInstancePools{
 			InstancePools: []compute.InstancePoolAndStats{
 				{
@@ -94,7 +47,7 @@ func AddHandlers(server *TestServer) {
 		}, nil
 	})
 
-	server.Handle("/api/2.1/clusters/list", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.1/clusters/list", func(r *http.Request) (any, error) {
 		return compute.ListClustersResponse{
 			Clusters: []compute.ClusterDetails{
 				{
@@ -109,13 +62,13 @@ func AddHandlers(server *TestServer) {
 		}, nil
 	})
 
-	server.Handle("/api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/preview/scim/v2/Me", func(r *http.Request) (any, error) {
 		return iam.User{
 			UserName: "tester@databricks.com",
 		}, nil
 	})
 
-	server.Handle("/api/2.0/workspace/get-status", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/workspace/get-status", func(r *http.Request) (any, error) {
 		return workspace.ObjectInfo{
 			ObjectId:   1001,
 			ObjectType: "DIRECTORY",
@@ -124,13 +77,13 @@ func AddHandlers(server *TestServer) {
 		}, nil
 	})
 
-	server.Handle("/api/2.1/unity-catalog/current-metastore-assignment", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.1/unity-catalog/current-metastore-assignment", func(r *http.Request) (any, error) {
 		return catalog.MetastoreAssignment{
 			DefaultCatalogName: "main",
 		}, nil
 	})
 
-	server.Handle("/api/2.0/permissions/directories/1001", func(r *http.Request) (any, error) {
+	server.Handle("GET /api/2.0/permissions/directories/1001", func(r *http.Request) (any, error) {
 		return workspace.WorkspaceObjectPermissions{
 			ObjectId:   "1001",
 			ObjectType: "DIRECTORY",

acceptance/workspace/jobs/create/output.txt

@@ -0,0 +1,5 @@
+
+>>> $CLI jobs create --json {"name":"abc"}
+{
+  "job_id":1111
+}

acceptance/workspace/jobs/create/script

@@ -0,0 +1 @@
+trace $CLI jobs create --json '{"name":"abc"}'

acceptance/workspace/jobs/create/server.json

@@ -0,0 +1,20 @@
+[
+  {
+    "method": "POST",
+    "path": "/api/2.1/jobs/create",
+    "request": {
+      "headers": {
+        "Authorization": "Bearer acceptance-test-token",
+        "User-Agent": "cli/[SEMVER] databricks-sdk-go/[SEMVER] go/[SEMVER] os/[OS] cmd/jobs_create cmd-exec-id/[UUID] auth/pat"
+      },
+      "body": {
+        "name": "abc"
+      }
+    },
+    "response": {
+      "body": {
+        "job_id": 1111
+      }
+    }
+  }
+]

cmd/root/user_agent_command_exec_id.go

@@ -7,8 +7,14 @@ import (
 	"github.com/google/uuid"
 )
 
+var cmdExecId = uuid.New().String()
+
+func CmdExecId() string {
+	return cmdExecId
+}
+
 func withCommandExecIdInUserAgent(ctx context.Context) context.Context {
 	// A UUID that will allow us to correlate multiple API requests made by
 	// the same CLI invocation.
-	return useragent.InContext(ctx, "cmd-exec-id", uuid.New().String())
+	return useragent.InContext(ctx, "cmd-exec-id", cmdExecId)
 }

integration/bundle/testdata/default_python/bundle_deploy.txt

@@ -1,5 +1,5 @@
 Building project_name_$UNIQUE_PRJ...
-Uploading project_name_$UNIQUE_PRJ-0.0.1+<NUMID>.<NUMID>-py3-none-any.whl...
+Uploading project_name_$UNIQUE_PRJ-0.0.1+[NUMID].[NUMID]-py3-none-any.whl...
 Uploading bundle files to /Workspace/Users/$USERNAME/.bundle/project_name_$UNIQUE_PRJ/dev/files...
 Deploying resources...
 Updating deployment state...

integration/bundle/testdata/default_python/bundle_summary.txt

@@ -16,7 +16,7 @@
         "enabled": false
       }
     },
-    "uuid": "<UUID>"
+    "uuid": "[UUID]"
   },
   "include": [
     "resources/project_name_$UNIQUE_PRJ.job.yml",
@@ -74,7 +74,7 @@
           ]
         },
         "format": "MULTI_TASK",
-        "id": "<NUMID>",
+        "id": "[NUMID]",
         "job_clusters": [
           {
             "job_cluster_key": "job_cluster",
@@ -141,7 +141,7 @@
             "unit": "DAYS"
           }
         },
-        "url": "$DATABRICKS_URL/jobs/<NUMID>?o=<NUMID>"
+        "url": "$DATABRICKS_URL/jobs/[NUMID]?o=[NUMID]"
       }
     },
     "pipelines": {
@@ -155,7 +155,7 @@
           "metadata_file_path": "/Workspace/Users/$USERNAME/.bundle/project_name_$UNIQUE_PRJ/dev/state/metadata.json"
         },
         "development": true,
-        "id": "<UUID>",
+        "id": "[UUID]",
         "libraries": [
           {
             "notebook": {
@@ -165,7 +165,7 @@
         ],
         "name": "[dev $USERNAME] project_name_$UNIQUE_PRJ_pipeline",
         "target": "project_name_$UNIQUE_PRJ_dev",
-        "url": "$DATABRICKS_URL/pipelines/<UUID>?o=<NUMID>"
+        "url": "$DATABRICKS_URL/pipelines/[UUID]?o=[NUMID]"
       }
     }
   },

internal/testserver/server.go

@@ -0,0 +1,135 @@
+package testserver
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/testdiff"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type Server struct {
+	*httptest.Server
+	Mux *http.ServeMux
+
+	t   testutil.TestingT
+	ctx context.Context
+
+	// API calls that we expect to be made.
+	calledPatterns map[string]bool
+}
+
+type ApiSpec struct {
+	Method  string `json:"method"`
+	Path    string `json:"path"`
+	Request struct {
+		Headers map[string]string `json:"headers"`
+		Body    json.RawMessage   `json:"body"`
+	} `json:"request"`
+	Response struct {
+		Body json.RawMessage `json:"body"`
+	} `json:"response"`
+}
+
+func New(ctx context.Context, t testutil.TestingT) *Server {
+	mux := http.NewServeMux()
+	server := httptest.NewServer(mux)
+
+	return &Server{
+		Server:         server,
+		Mux:            mux,
+		t:              t,
+		ctx:            ctx,
+		calledPatterns: make(map[string]bool),
+	}
+}
+
+func NewFromConfig(t testutil.TestingT, path string) *Server {
+	content := testutil.ReadFile(t, path)
+	var apiSpecs []ApiSpec
+	err := json.Unmarshal([]byte(content), &apiSpecs)
+	require.NoError(t, err)
+
+	ctx, replacements := testdiff.WithReplacementsMap(context.Background())
+	testdiff.PrepareReplacementOS(t, replacements)
+	testdiff.PrepareReplacementsUUID(t, replacements)
+	testdiff.PrepareReplacementsSemver(t, replacements)
+
+	server := New(ctx, t)
+	for _, apiSpec := range apiSpecs {
+		server.MustHandle(apiSpec)
+	}
+
+	return server
+}
+
+type HandlerFunc func(req *http.Request) (resp any, err error)
+
+func (s *Server) MustHandle(apiSpec ApiSpec) {
+	require.NotEmpty(s.t, apiSpec.Method)
+	require.NotEmpty(s.t, apiSpec.Path)
+
+	pattern := apiSpec.Method + " " + apiSpec.Path
+	s.calledPatterns[pattern] = false
+
+	s.Handle(pattern, func(req *http.Request) (any, error) {
+		for k, v := range apiSpec.Request.Headers {
+			testdiff.AssertEqualStrings(s.t, s.ctx, v, req.Header.Get(k))
+		}
+
+		b, err := io.ReadAll(req.Body)
+		require.NoError(s.t, err)
+
+		// Assert that the request body matches the expected body.
+		assert.JSONEq(s.t, string(apiSpec.Request.Body), string(b))
+
+		// Record the fact that this pattern was called.
+		s.calledPatterns[pattern] = true
+
+		// Return the expected response body.
+		return apiSpec.Response.Body, err
+	})
+}
+
+func (s *Server) Close() {
+	for pattern, called := range s.calledPatterns {
+		assert.Truef(s.t, called, "expected pattern %s to be called", pattern)
+	}
+
+	s.Server.Close()
+}
+
+func (s *Server) Handle(pattern string, handler HandlerFunc) {
+	s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
+		resp, err := handler(r)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+
+		var respBytes []byte
+
+		respString, ok := resp.(string)
+		if ok {
+			respBytes = []byte(respString)
+		} else {
+			respBytes, err = json.MarshalIndent(resp, "", "    ")
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+
+		if _, err := w.Write(respBytes); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+	})
+}

internal/testutil/file.go

@@ -62,6 +62,18 @@ func StatFile(t TestingT, path string) os.FileInfo {
 	return fi
 }
 
+func DetectFile(t TestingT, path string) bool {
+	_, err := os.Stat(path)
+	if err == nil {
+		return true
+	}
+	if os.IsNotExist(err) {
+		return false
+	}
+	t.Fatalf("unexpected error: %v", err)
+	return false
+}
+
 // AssertFileContents asserts that the file at path has the expected content.
 func AssertFileContents(t TestingT, path, expected string) bool {
 	actual := ReadFile(t, path)

libs/testdiff/golden.go

@@ -11,7 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-var OverwriteMode = false
+var OverwriteMode = true
 
 func init() {
 	flag.BoolVar(&OverwriteMode, "update", false, "Overwrite golden files")
@@ -65,6 +65,16 @@ func AssertOutputJQ(t testutil.TestingT, ctx context.Context, out, outTitle, exp
 	}
 }
 
+func AssertEqualStrings(t testutil.TestingT, ctx context.Context, expected, actual string) {
+	t.Helper()
+	replacements := GetReplacementsMap(ctx)
+	if replacements == nil {
+		t.Fatal("WithReplacementsMap was not called")
+	}
+	actual = replacements.Replace(actual)
+	assert.Equal(t, expected, actual)
+}
+
 func ReplaceOutput(t testutil.TestingT, ctx context.Context, out string) string {
 	t.Helper()
 	out = NormalizeNewlines(out)

libs/testdiff/replacement.go

@@ -20,6 +20,8 @@ const (
 )
 
 var (
+	// From https://semver.org/#is-there-a-suggested-regular-expression-regex-to-check-a-semver-string
+	semverRegex      = regexp.MustCompile(`(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?`)
 	uuidRegex        = regexp.MustCompile(`[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}`)
 	numIdRegex       = regexp.MustCompile(`[0-9]{3,}`)
 	privatePathRegex = regexp.MustCompile(`(/tmp|/private)(/.*)/([a-zA-Z0-9]+)`)
@@ -183,15 +185,25 @@ func PrepareReplacementsUser(t testutil.TestingT, r *ReplacementsContext, u iam.
 
 func PrepareReplacementsUUID(t testutil.TestingT, r *ReplacementsContext) {
 	t.Helper()
-	r.append(uuidRegex, "<UUID>")
+	r.append(uuidRegex, "[UUID]")
 }
 
 func PrepareReplacementsNumber(t testutil.TestingT, r *ReplacementsContext) {
 	t.Helper()
-	r.append(numIdRegex, "<NUMID>")
+	r.append(numIdRegex, "[NUMID]")
 }
 
 func PrepareReplacementsTemporaryDirectory(t testutil.TestingT, r *ReplacementsContext) {
 	t.Helper()
 	r.append(privatePathRegex, "/tmp/.../$3")
 }
+
+func PrepareReplacementsSemver(t testutil.TestingT, r *ReplacementsContext) {
+	t.Helper()
+	r.append(semverRegex, "[SEMVER]")
+}
+
+func PrepareReplacementOS(t testutil.TestingT, r *ReplacementsContext) {
+	t.Helper()
+	r.Set(runtime.GOOS, "[OS]")
+}

libs/testdiff/replacement_test.go

@@ -25,7 +25,7 @@ func TestReplacement_UUID(t *testing.T) {
 
 	PrepareReplacementsUUID(t, &repls)
 
-	assert.Equal(t, "<UUID>", repls.Replace("123e4567-e89b-12d3-a456-426614174000"))
+	assert.Equal(t, "[UUID]", repls.Replace("123e4567-e89b-12d3-a456-426614174000"))
 }
 
 func TestReplacement_Number(t *testing.T) {
@@ -34,7 +34,7 @@ func TestReplacement_Number(t *testing.T) {
 	PrepareReplacementsNumber(t, &repls)
 
 	assert.Equal(t, "12", repls.Replace("12"))
-	assert.Equal(t, "<NUMID>", repls.Replace("123"))
+	assert.Equal(t, "[NUMID]", repls.Replace("123"))
 }
 
 func TestReplacement_TemporaryDirectory(t *testing.T) {