Add feature to mock server APIs in acceptance tests

acceptance/acceptance_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/testdiff"
+	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/stretchr/testify/require"
 )
@@ -109,10 +110,10 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	cloudEnv := os.Getenv("CLOUD_ENV")
 
 	if cloudEnv == "" {
-		server := StartServer(t)
-		AddHandlers(server)
+		defaultServer := StartServer(t)
+		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")
 
 		homeDir := t.TempDir()
@@ -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()
+		})
+	}
+
 	repls.SetPathWithParents(tmpDir, "$TMPDIR")
 
 	scriptContents := readMergedScriptContents(t, dir)

acceptance/cmd_server_test.go

@@ -8,10 +8,11 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/internal/testcli"
+	"github.com/databricks/cli/libs/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,73 +1,25 @@
 package acceptance_test
 
 import (
-	"encoding/json"
 	"net/http"
-	"net/http/httptest"
 	"testing"
 
+	"github.com/databricks/cli/libs/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 {
-	*httptest.Server
-	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()
+func StartServer(t *testing.T) *testserver.Server {
+	server := testserver.New(t)
 	t.Cleanup(func() {
 		server.Close()
 	})
 	return server
 }
 
-func AddHandlers(server *TestServer) {
+func AddHandlers(server *testserver.Server) {
 	server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, error) {
 		return compute.ListPoliciesResponse{
 			Policies: []compute.Policy{

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,11 @@
+[
+  {
+    "method": "POST",
+    "path": "/api/2.1/jobs/create",
+    "response": {
+      "body": {
+        "job_id": 1111
+      }
+    }
+  }
+]

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/testserver/server.go

@@ -0,0 +1,111 @@
+package testserver
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type Server struct {
+	*httptest.Server
+	Mux *http.ServeMux
+
+	t testutil.TestingT
+
+	// API calls that we expect to be made.
+	calledPatterns map[string]bool
+}
+
+type ApiSpec struct {
+	Method   string `json:"method"`
+	Path     string `json:"path"`
+	Response struct {
+		Body json.RawMessage `json:"body"`
+	} `json:"response"`
+}
+
+func New(t testutil.TestingT) *Server {
+	mux := http.NewServeMux()
+	server := httptest.NewServer(mux)
+
+	return &Server{
+		Server:         server,
+		Mux:            mux,
+		t:              t,
+		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)
+
+	server := New(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) {
+	assert.NotEmpty(s.t, apiSpec.Method)
+	assert.NotEmpty(s.t, apiSpec.Path)
+
+	pattern := apiSpec.Method + " " + apiSpec.Path
+	s.calledPatterns[pattern] = false
+
+	s.Handle(pattern, func(req *http.Request) (any, error) {
+		// Record the fact that this pattern was called.
+		s.calledPatterns[pattern] = true
+
+		// Return the expected response body.
+		return apiSpec.Response.Body, nil
+	})
+}
+
+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
+		}
+	})
+}