From 57b8d336e03cb2b9aa0c353d9cce7d85d5c72c7f Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Wed, 5 Feb 2025 15:02:15 +0530 Subject: [PATCH] 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 | 3 +++ acceptance/config_test.go | 3 +++ .../workspace/jobs/create/out.requests.txt | 2 +- acceptance/workspace/jobs/create/test.toml | 17 ++++++++++++ libs/testdiff/replacement.go | 18 +++++++++++++ libs/testserver/server.go | 26 ++++++++++++++----- 6 files changed, 61 insertions(+), 8 deletions(-) diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 6e302fa96..f205217ff 100644 --- a/acceptance/acceptance_test.go +++ b/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 { diff --git a/acceptance/config_test.go b/acceptance/config_test.go index 97be457e2..e24a683e7 100644 --- a/acceptance/config_test.go +++ b/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 { diff --git a/acceptance/workspace/jobs/create/out.requests.txt b/acceptance/workspace/jobs/create/out.requests.txt index b22876b70..4a85c4c43 100644 --- a/acceptance/workspace/jobs/create/out.requests.txt +++ b/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"}} diff --git a/acceptance/workspace/jobs/create/test.toml b/acceptance/workspace/jobs/create/test.toml index e69569c18..1fd9b3cec 100644 --- a/acceptance/workspace/jobs/create/test.toml +++ b/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 = "" diff --git a/libs/testdiff/replacement.go b/libs/testdiff/replacement.go index 5bbba1be1..d4d5eb27b 100644 --- a/libs/testdiff/replacement.go +++ b/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]") +} diff --git a/libs/testserver/server.go b/libs/testserver/server.go index a751531ed..5e3efe1c5 100644 --- a/libs/testserver/server.go +++ b/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"` - Path string `json:"path"` - Body any `json:"body"` + Headers map[string]string `json:"headers,omitempty"` + Method string `json:"method"` + 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, - Path: r.URL.Path, - Body: json.RawMessage(body), + Headers: headers, + Method: r.Method, + Path: r.URL.Path, + Body: json.RawMessage(body), }) }