Add support for recording requests in integration acceptance tests

.wsignore

@@ -11,6 +11,12 @@ experimental/python/docs/images/databricks-logo.svg
 **/*.zip
 **/*.whl
 
+# new lines are recorded differently on windows and unix.
+# In unix: "raw_body": "hello, world\n"
+# In windows: "raw_body": "hello, world\r\n"
+# In order to prevent that difference, hello.txt does not have a trailing newline.
+acceptance/selftest/record_cloud/volume-io/hello.txt
+
 # "bundle run" has trailing whitespace:
 acceptance/bundle/integration_whl/*/output.txt
 

acceptance/acceptance_test.go

@@ -18,7 +18,6 @@ import (
 	"slices"
 	"sort"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 	"unicode/utf8"
@@ -242,15 +241,15 @@ func testAccept(t *testing.T, inprocessMode bool, singleTest string) int {
 				if len(expanded[0]) > 0 {
 					t.Logf("Running test with env %v", expanded[0])
 				}
-				runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0])
+				runTest(t, dir, coverDir, repls.Clone(), config, configPath, expanded[0], inprocessMode)
 			} else {
 				for _, envset := range expanded {
 					envname := strings.Join(envset, "/")
 					t.Run(envname, func(t *testing.T) {
 						if !inprocessMode {
 							t.Parallel()
 						}
-						runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset)
+						runTest(t, dir, coverDir, repls.Clone(), config, configPath, envset, inprocessMode)
 					})
 				}
 			}
@@ -342,7 +341,14 @@ func getSkipReason(config *internal.TestConfig, configPath string) string {
 	return ""
 }
 
-func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsContext, config internal.TestConfig, configPath string, customEnv []string) {
+func runTest(t *testing.T,
+	dir, coverDir string,
+	repls testdiff.ReplacementsContext,
+	config internal.TestConfig,
+	configPath string,
+	customEnv []string,
+	inprocessMode bool,
+) {
 	if LogConfig {
 		configBytes, err := json.MarshalIndent(config, "", "  ")
 		require.NoError(t, err)
@@ -395,16 +401,12 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	} else if isRunningOnCloud {
 		timeout = max(timeout, config.TimeoutCloud)
 	}
-
 	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
 	defer cancelFunc()
 	args := []string{"bash", "-euo", "pipefail", EntryPointScript}
 	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
 
-	// This mutex is used to synchronize recording requests
-	var serverMutex sync.Mutex
-
-	cfg, user := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir, &serverMutex)
+	cfg, user := internal.PrepareServerAndClient(t, config, LogRequests, tmpDir)
 	testdiff.PrepareReplacementsUser(t, &repls, user)
 	testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg)
 

acceptance/internal/prepare_server.go

@@ -14,12 +14,12 @@ import (
 	"time"
 	"unicode/utf8"
 
-	sdkconfig "github.com/databricks/databricks-sdk-go/config"
-	"github.com/databricks/databricks-sdk-go/service/iam"
-
 	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/testproxy"
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
+	sdkconfig "github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -40,81 +40,115 @@ func isTruePtr(value *bool) bool {
 	return value != nil && *value
 }
 
-func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string, mu *sync.Mutex) (*sdkconfig.Config, iam.User) {
+func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, outputDir string) (*sdkconfig.Config, iam.User) {
 	cloudEnv := os.Getenv("CLOUD_ENV")
+	recordRequests := isTruePtr(config.RecordRequests)
 
-	// If we are running on a cloud environment, use the host configured in the
-	// environment.
 	if cloudEnv != "" {
 		w, err := databricks.NewWorkspaceClient()
 		require.NoError(t, err)
 
 		user, err := w.CurrentUser.Me(context.Background())
 		require.NoError(t, err, "Failed to get current user")
 
-		return w.Config, *user
-	}
+		cfg := w.Config
 
-	recordRequests := isTruePtr(config.RecordRequests)
+		// If we are running in a cloud environment AND we are recording requests,
+		// start a dedicated server to act as a reverse proxy to a real Databricks workspace.
+		if recordRequests {
+			host, token := startProxyServer(t, logRequests, config.IncludeRequestHeaders, outputDir)
+			cfg = &sdkconfig.Config{
+				Host:  host,
+				Token: token,
+			}
+		}
 
-	tokenSuffix := strings.ReplaceAll(uuid.NewString(), "-", "")
-	token := "dbapi" + tokenSuffix
+		return cfg, *user
+	}
 
-	// If we are not recording requests, and no custom server server stubs are configured,
+	// If we are not recording requests, and no custom server stubs are configured,
 	// use the default shared server.
 	if len(config.Server) == 0 && !recordRequests {
-		return &sdkconfig.Config{
+		// 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
+
+		cfg := &sdkconfig.Config{
 			Host:  os.Getenv("DATABRICKS_DEFAULT_HOST"),
 			Token: token,
-		}, TestUser
-	}
+		}
 
-	host := startDedicatedServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir, mu)
+		return cfg, TestUser
+	}
 
-	return &sdkconfig.Config{
+	// Default case. Start a dedicated local server for the test with the server stubs configured
+	// as overrides.
+	host, token := startLocalServer(t, config.Server, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir)
+	cfg := &sdkconfig.Config{
 		Host:  host,
 		Token: token,
-	}, TestUser
+	}
+
+	// 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
 }
 
-func startDedicatedServer(t *testing.T,
+func recordRequestsCallback(t *testing.T, includeHeaders []string, outputDir string) func(request *testserver.Request) {
+	mu := sync.Mutex{}
+
+	return func(request *testserver.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		req := getLoggedRequest(request, includeHeaders)
+		reqJson, err := json.MarshalIndent(req, "", "  ")
+		assert.NoErrorf(t, err, "Failed to json-encode: %#v", req)
+
+		requestsPath := filepath.Join(outputDir, "out.requests.txt")
+		f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
+		assert.NoError(t, err)
+		defer f.Close()
+
+		_, err = f.WriteString(string(reqJson) + "\n")
+		assert.NoError(t, err)
+	}
+}
+
+func logResponseCallback(t *testing.T) func(request *testserver.Request, response *testserver.EncodedResponse) {
+	mu := sync.Mutex{}
+
+	return func(request *testserver.Request, response *testserver.EncodedResponse) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		t.Logf("%d %s %s\n%s\n%s",
+			response.StatusCode, request.Method, request.URL,
+			formatHeadersAndBody("> ", request.Headers, request.Body),
+			formatHeadersAndBody("# ", response.Headers, response.Body),
+		)
+	}
+}
+
+func startLocalServer(t *testing.T,
 	stubs []ServerStub,
 	recordRequests bool,
 	logRequests bool,
 	includeHeaders []string,
 	outputDir string,
-	mu *sync.Mutex,
-) string {
+) (string, string) {
 	s := testserver.New(t)
 
+	// Record API requests in out.requests.txt if RecordRequests is true
+	// in test.toml
 	if recordRequests {
-		requestsPath := filepath.Join(outputDir, "out.requests.txt")
-		s.RequestCallback = func(request *testserver.Request) {
-			req := getLoggedRequest(request, includeHeaders)
-			reqJson, err := json.MarshalIndent(req, "", "  ")
-
-			mu.Lock()
-			defer mu.Unlock()
-
-			assert.NoErrorf(t, err, "Failed to json-encode: %#v", req)
-
-			f, err := os.OpenFile(requestsPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
-			assert.NoError(t, err)
-			defer f.Close()
-
-			_, err = f.WriteString(string(reqJson) + "\n")
-			assert.NoError(t, err)
-		}
+		s.RequestCallback = recordRequestsCallback(t, includeHeaders, outputDir)
 	}
 
+	// Log API responses if the -logrequests flag is set.
 	if logRequests {
-		s.ResponseCallback = func(request *testserver.Request, response *testserver.EncodedResponse) {
-			t.Logf("%d %s %s\n%s\n%s",
-				response.StatusCode, request.Method, request.URL,
-				formatHeadersAndBody("> ", request.Headers, request.Body),
-				formatHeadersAndBody("# ", response.Headers, response.Body),
-			)
-		}
+		s.ResponseCallback = logResponseCallback(t)
 	}
 
 	for ind := range stubs {
@@ -132,8 +166,25 @@ func startDedicatedServer(t *testing.T,
 
 	// The earliest handlers take precedence, add default handlers last
 	addDefaultHandlers(s)
+	return s.URL, "dbapi123"
+}
+
+func startProxyServer(t *testing.T,
+	logRequests bool,
+	includeHeaders []string,
+	outputDir string,
+) (string, string) {
+	s := testproxy.New(t)
+
+	// Always record requests for a proxy server.
+	s.RequestCallback = recordRequestsCallback(t, includeHeaders, outputDir)
+
+	// Log API responses if the -logrequests flag is set.
+	if logRequests {
+		s.ResponseCallback = logResponseCallback(t)
+	}
 
-	return s.URL
+	return s.URL, "dbapi1234"
 }
 
 type LoggedRequest struct {

acceptance/selftest/record_cloud/basic/out.requests.txt

@@ -0,0 +1,4 @@
+{
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}

acceptance/selftest/record_cloud/basic/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] current-user me
+"[USERNAME]"

acceptance/selftest/record_cloud/basic/script

@@ -0,0 +1,2 @@
+# Proxy server successfully records a requests and returns a response.
+trace $CLI current-user me | jq .name.givenName

acceptance/selftest/record_cloud/error/out.requests.txt

@@ -0,0 +1,4 @@
+{
+  "method": "GET",
+  "path": "/api/2.2/jobs/get"
+}

acceptance/selftest/record_cloud/error/output.txt

@@ -0,0 +1,5 @@
+
+>>> [CLI] jobs get 1234
+Error: Job 1234 does not exist.
+
+Exit code: 1

acceptance/selftest/record_cloud/error/script

@@ -0,0 +1,2 @@
+# Proxy server should successfully return non 200 responses.
+errcode trace $CLI jobs get 1234

acceptance/selftest/record_cloud/pipeline-crud/output.txt

@@ -0,0 +1,80 @@
+
+=== Create a pipeline
+
+>>> print_requests
+{
+  "method": "POST",
+  "path": "/api/2.0/pipelines",
+  "body": {
+    "allow_duplicate_names": true,
+    "libraries": [
+      {
+        "file": {
+          "path": "/whatever.py"
+        }
+      }
+    ],
+    "name": "test-pipeline-1"
+  }
+}
+
+=== Get the pipeline
+>>> [CLI] pipelines get [UUID]
+"test-pipeline-1"
+
+>>> print_requests
+{
+  "method": "GET",
+  "path": "/api/2.0/pipelines/[UUID]"
+}
+
+=== Update the pipeline
+>>> [CLI] pipelines update [UUID] --json @pipeline2.json
+
+>>> print_requests
+{
+  "method": "PUT",
+  "path": "/api/2.0/pipelines/[UUID]",
+  "body": {
+    "allow_duplicate_names": true,
+    "libraries": [
+      {
+        "file": {
+          "path": "/whatever.py"
+        }
+      }
+    ],
+    "name": "test-pipeline-2"
+  }
+}
+
+=== Verify the update
+>>> [CLI] pipelines get [UUID]
+"test-pipeline-2"
+
+>>> print_requests
+{
+  "method": "GET",
+  "path": "/api/2.0/pipelines/[UUID]"
+}
+
+=== Delete the pipeline
+>>> [CLI] pipelines delete [UUID]
+
+>>> print_requests
+{
+  "method": "DELETE",
+  "path": "/api/2.0/pipelines/[UUID]"
+}
+
+=== Verify the deletion
+>>> [CLI] pipelines get [UUID]
+Error: The specified pipeline [UUID] was not found.
+
+Exit code: 1
+
+>>> print_requests
+{
+  "method": "GET",
+  "path": "/api/2.0/pipelines/[UUID]"
+}

acceptance/selftest/record_cloud/pipeline-crud/pipeline1.json

@@ -0,0 +1,11 @@
+{
+    "name": "test-pipeline-1",
+    "allow_duplicate_names": true,
+    "libraries": [
+        {
+            "file": {
+                "path": "/whatever.py"
+            }
+        }
+    ]
+}