Add ability to record headers in acceptance tests (#2296)

## Changes
HTTP headers like the User-Agent are an important part of our internal
ETL pipelines. This PR adds the ability to validate the headers used in
an HTTP request as part of our acceptance tests.

## Tests
Modifying existing test.

acceptance/acceptance_test.go

@@ -156,6 +156,8 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int {
 	testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
 	testdiff.PrepareReplacementsUUID(t, &repls)
 	testdiff.PrepareReplacementsDevVersion(t, &repls)
+	testdiff.PrepareReplacementSdkVersion(t, &repls)
+	testdiff.PrepareReplacementsGoVersion(t, &repls)
 
 	testDirs := getTests(t)
 	require.NotEmpty(t, testDirs)
@@ -253,6 +255,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
 	if len(config.Server) > 0 || config.RecordRequests {
 		server = testserver.New(t)
 		server.RecordRequests = config.RecordRequests
+		server.IncludeRequestHeaders = config.IncludeRequestHeaders
 
 		// If no custom server stubs are defined, add the default handlers.
 		if len(config.Server) == 0 {

acceptance/config_test.go

@@ -47,6 +47,9 @@ type TestConfig struct {
 	// Record the requests made to the server and write them as output to
 	// out.requests.txt
 	RecordRequests bool
+
+	// List of request headers to include when recording requests.
+	IncludeRequestHeaders []string
 }
 
 type ServerStub struct {

acceptance/workspace/jobs/create/out.requests.txt

@@ -1 +1 @@
-{"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}}
+{"headers":{"Authorization":"Bearer dapi1234","User-Agent":"cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/jobs_create cmd-exec-id/[UUID] auth/pat"},"method":"POST","path":"/api/2.1/jobs/create","body":{"name":"abc"}}

acceptance/workspace/jobs/create/test.toml

@@ -1,5 +1,6 @@
 LocalOnly = true  # request recording currently does not work with cloud environment
 RecordRequests = true
+IncludeRequestHeaders = ["Authorization", "User-Agent"]
 
 [[Server]]
 Pattern = "POST /api/2.1/jobs/create"
@@ -8,3 +9,19 @@ Response.Body = '''
     "job_id": 1111
 }
 '''
+
+[[Repls]]
+Old = "(linux|darwin|windows)"
+New = "[OS]"
+
+[[Repls]]
+Old = " upstream/[A-Za-z0-9.-]+"
+New = ""
+
+[[Repls]]
+Old = " upstream-version/[A-Za-z0-9.-]+"
+New = ""
+
+[[Repls]]
+Old = " cicd/[A-Za-z0-9.-]+"
+New = ""

libs/testdiff/replacement.go

@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/libs/iamutil"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/iam"
+	"golang.org/x/mod/semver"
 )
 
 const (
@@ -208,3 +209,20 @@ func PrepareReplacementsDevVersion(t testutil.TestingT, r *ReplacementsContext)
 	t.Helper()
 	r.append(devVersionRegex, "[DEV_VERSION]")
 }
+
+func PrepareReplacementSdkVersion(t testutil.TestingT, r *ReplacementsContext) {
+	t.Helper()
+	r.Set(databricks.Version(), "[SDK_VERSION]")
+}
+
+func goVersion() string {
+	gv := runtime.Version()
+	ssv := strings.ReplaceAll(gv, "go", "v")
+	sv := semver.Canonical(ssv)
+	return strings.TrimPrefix(sv, "v")
+}
+
+func PrepareReplacementsGoVersion(t testutil.TestingT, r *ReplacementsContext) {
+	t.Helper()
+	r.Set(goVersion(), "[GO_VERSION]")
+}

libs/testserver/server.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"slices"
 
 	"github.com/stretchr/testify/assert"
 
@@ -17,15 +18,17 @@ type Server struct {
 
 	t testutil.TestingT
 
-	RecordRequests bool
+	RecordRequests        bool
+	IncludeRequestHeaders []string
 
 	Requests []Request
 }
 
 type Request struct {
-	Method string `json:"method"`
+	Headers map[string]string `json:"headers,omitempty"`
-	Path   string `json:"path"`
+	Method  string            `json:"method"`
-	Body   any    `json:"body"`
+	Path    string            `json:"path"`
+	Body    any               `json:"body"`
 }
 
 func New(t testutil.TestingT) *Server {
@@ -50,10 +53,19 @@ func (s *Server) Handle(pattern string, handler HandlerFunc) {
 			body, err := io.ReadAll(r.Body)
 			assert.NoError(s.t, err)
 
+			headers := make(map[string]string)
+			for k, v := range r.Header {
+				if len(v) == 0 || !slices.Contains(s.IncludeRequestHeaders, k) {
+					continue
+				}
+				headers[k] = v[0]
+			}
+
 			s.Requests = append(s.Requests, Request{
-				Method: r.Method,
+				Headers: headers,
-				Path:   r.URL.Path,
+				Method:  r.Method,
-				Body:   json.RawMessage(body),
+				Path:    r.URL.Path,
+				Body:    json.RawMessage(body),
 			})
 
 		}