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.
This commit is contained in:
shreyas-goenka 2025-02-05 15:02:15 +05:30 committed by GitHub
parent 1678503cb0
commit 57b8d336e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 61 additions and 8 deletions

View File

@ -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 {

View File

@ -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 {

View File

@ -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"}}

View File

@ -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 = ""

View File

@ -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]")
}

View File

@ -5,6 +5,7 @@ import (
"io"
"net/http"
"net/http/httptest"
"slices"
"github.com/stretchr/testify/assert"
@ -18,11 +19,13 @@ type Server struct {
t testutil.TestingT
RecordRequests bool
IncludeRequestHeaders []string
Requests []Request
}
type Request struct {
Headers map[string]string `json:"headers,omitempty"`
Method string `json:"method"`
Path string `json:"path"`
Body any `json:"body"`
@ -50,7 +53,16 @@ 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{
Headers: headers,
Method: r.Method,
Path: r.URL.Path,
Body: json.RawMessage(body),