From 0264c9e6b2351377ac6b04fd83c57e1b138871e9 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 28 Oct 2026 16:18:43 +0000 Subject: [PATCH 1/3] Extract shared fault rules into libs/fault --- libs/fault/fault.go | 132 ++++++++++++++++++++++++++++++++++ libs/fault/fault_test.go | 69 ++++++++++++++++++ libs/testserver/fault.go | 127 +------------------------------- libs/testserver/fault_test.go | 60 ---------------- libs/testserver/server.go | 7 +- 5 files changed, 208 insertions(+), 187 deletions(-) create mode 100644 libs/fault/fault.go create mode 100644 libs/fault/fault_test.go diff --git a/libs/fault/fault.go b/libs/fault/fault.go new file mode 100644 index 00000000000..9da0d579d72 --- /dev/null +++ b/libs/fault/fault.go @@ -0,0 +1,132 @@ +package fault + +import ( + "bytes" + "encoding/json" + "strings" + "sync" +) + +type faultRuleKey struct { + token string + pattern string +} + +// Rule describes a single injected fault: HTTP status, body, and remaining fire count. +type Rule struct { + StatusCode int + Body string + // AfterHandler keeps the handler's state change and replaces only its response. + AfterHandler bool + offset int + times int + + // bodyContains, if non-empty, additionally requires the request body to + // contain this substring for the rule to fire. /workspace/import routes + // every upload through the same method+path, so the only way to target a + // single file's upload is by matching its multipart "path" form field. + bodyContains string +} + +// Rules holds the active fault injection rules for a test server. +type Rules struct { + mu sync.Mutex + rules map[faultRuleKey]*Rule +} + +// NewRules returns an empty Rules. +func NewRules() *Rules { + return &Rules{rules: make(map[faultRuleKey]*Rule)} +} + +// Set registers or replaces a fault rule for the given token and pattern. +func (fr *Rules) Set(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { + fr.set(token, pattern, statusCode, body, offset, times, false, bodyContains) +} + +// SetAfterHandler is like Set, but the handler runs first so its effect is kept. +func (fr *Rules) SetAfterHandler(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { + fr.set(token, pattern, statusCode, body, offset, times, true, bodyContains) +} + +func (fr *Rules) set(token, pattern string, statusCode int, body string, offset, times int, afterHandler bool, bodyContains string) { + fr.mu.Lock() + defer fr.mu.Unlock() + fr.rules[faultRuleKey{token: token, pattern: pattern}] = &Rule{ + StatusCode: statusCode, + Body: body, + AfterHandler: afterHandler, + offset: offset, + times: times, + bodyContains: bodyContains, + } +} + +// Check returns a matching fault rule and advances its counters, or nil if no rule matches. +// Pattern supports a trailing "*" wildcard, e.g. "PUT /api/2.0/permissions/pipelines/*". +// A rule with a non-empty bodyContains only matches when body contains that substring. +func (fr *Rules) Check(method, path, token string, body []byte) *Rule { + requestPattern := method + " " + path + + fr.mu.Lock() + defer fr.mu.Unlock() + + for key, rule := range fr.rules { + if key.token != token { + continue + } + rulePattern := key.pattern + var matched bool + if strings.HasSuffix(rulePattern, "*") { + matched = strings.HasPrefix(requestPattern, rulePattern[:len(rulePattern)-1]) + } else { + matched = requestPattern == rulePattern + } + if !matched { + continue + } + // A non-matching body must not consume the rule's offset/times budget, + // so this check precedes the counter updates below. + if rule.bodyContains != "" && !bytes.Contains(body, []byte(rule.bodyContains)) { + continue + } + if rule.offset > 0 { + rule.offset-- + return nil + } + if rule.times <= 0 { + delete(fr.rules, key) + return nil + } + rule.times-- + if rule.times == 0 { + delete(fr.rules, key) + } + result := *rule + return &result + } + + return nil +} + +// Register decodes and registers a rule for token using the fault.py wire format. +func (fr *Rules) Register(token string, data []byte) error { + var body struct { + Pattern string `json:"pattern"` + StatusCode int `json:"status_code"` + Body string `json:"body"` + Offset int `json:"offset"` + Times int `json:"times"` + AfterHandler bool `json:"after_handler"` + BodyContains string `json:"body_contains"` + } + if err := json.Unmarshal(data, &body); err != nil { + return err + } + if body.AfterHandler { + fr.SetAfterHandler(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } else { + fr.Set(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } + return nil +} diff --git a/libs/fault/fault_test.go b/libs/fault/fault_test.go new file mode 100644 index 00000000000..202f7761a08 --- /dev/null +++ b/libs/fault/fault_test.go @@ -0,0 +1,69 @@ +package fault_test + +import ( + "testing" + + "github.com/databricks/cli/libs/fault" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestFaultRulesNoMatch(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "GET /foo", 504, "body", 0, 1, "") + + assert.Nil(t, fr.Check("POST", "/foo", "tok", nil)) + assert.Nil(t, fr.Check("GET", "/bar", "tok", nil)) + assert.Nil(t, fr.Check("GET", "/foo", "other", nil)) +} + +func TestFaultRulesExactMatch(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "PUT /api/2.0/jobs/123", 504, "body", 0, 1, "") + + rule := fr.Check("PUT", "/api/2.0/jobs/123", "tok", nil) + require.NotNil(t, rule) + assert.Equal(t, 504, rule.StatusCode) + assert.Equal(t, "body", rule.Body) +} + +func TestFaultRulesWildcardMatch(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "PUT /api/2.0/permissions/pipelines/*", 504, "body", 0, 2, "") + + assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/abc", "tok", nil)) + assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) + assert.Nil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) // exhausted +} + +func TestFaultRulesOffset(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "GET /foo", 504, "body", 2, 1, "") + + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 2→1 + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 1→0 + assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) // fires + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted +} + +func TestFaultRulesTimes(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "GET /foo", 504, "body", 0, 3, "") + + for range 3 { + assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) + } + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted +} + +func TestFaultRulesBodyContains(t *testing.T) { + fr := fault.NewRules() + fr.Set("tok", "POST /api/2.0/workspace/import", 403, "body", 0, 1, "state/resources.json") + + // Body without the substring must not match and must not consume the budget. + assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/deploy.lock"))) + // Body with the substring fires. + assert.NotNil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) + // Budget was only consumed by the matching request. + assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) +} diff --git a/libs/testserver/fault.go b/libs/testserver/fault.go index 43d46f3f28c..fff8db13a8a 100644 --- a/libs/testserver/fault.go +++ b/libs/testserver/fault.go @@ -1,134 +1,13 @@ package testserver -import ( - "bytes" - "encoding/json" - "strings" - "sync" -) - -type faultRuleKey struct { - token string - pattern string -} - -// FaultRule describes a single injected fault: HTTP status, body, and remaining fire count. -type FaultRule struct { - StatusCode int - Body string - // AfterHandler keeps the handler's state change and replaces only its response. - AfterHandler bool - offset int - times int - - // bodyContains, if non-empty, additionally requires the request body to - // contain this substring for the rule to fire. /workspace/import routes - // every upload through the same method+path, so the only way to target a - // single file's upload is by matching its multipart "path" form field. - bodyContains string -} - -// FaultRules holds the active fault injection rules for a test server. -type FaultRules struct { - mu sync.Mutex - rules map[faultRuleKey]*FaultRule -} - -// NewFaultRules returns an empty FaultRules. -func NewFaultRules() *FaultRules { - return &FaultRules{rules: make(map[faultRuleKey]*FaultRule)} -} - -// Set registers or replaces a fault rule for the given token and pattern. -func (fr *FaultRules) Set(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { - fr.set(token, pattern, statusCode, body, offset, times, false, bodyContains) -} - -// SetAfterHandler is like Set, but the handler runs first so its effect is kept. -func (fr *FaultRules) SetAfterHandler(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { - fr.set(token, pattern, statusCode, body, offset, times, true, bodyContains) -} - -func (fr *FaultRules) set(token, pattern string, statusCode int, body string, offset, times int, afterHandler bool, bodyContains string) { - fr.mu.Lock() - defer fr.mu.Unlock() - fr.rules[faultRuleKey{token: token, pattern: pattern}] = &FaultRule{ - StatusCode: statusCode, - Body: body, - AfterHandler: afterHandler, - offset: offset, - times: times, - bodyContains: bodyContains, - } -} - -// Check returns a matching fault rule and advances its counters, or nil if no rule matches. -// Pattern supports a trailing "*" wildcard, e.g. "PUT /api/2.0/permissions/pipelines/*". -// A rule with a non-empty bodyContains only matches when body contains that substring. -func (fr *FaultRules) Check(method, path, token string, body []byte) *FaultRule { - requestPattern := method + " " + path - - fr.mu.Lock() - defer fr.mu.Unlock() - - for key, rule := range fr.rules { - if key.token != token { - continue - } - rulePattern := key.pattern - var matched bool - if strings.HasSuffix(rulePattern, "*") { - matched = strings.HasPrefix(requestPattern, rulePattern[:len(rulePattern)-1]) - } else { - matched = requestPattern == rulePattern - } - if !matched { - continue - } - // A non-matching body must not consume the rule's offset/times budget, - // so this check precedes the counter updates below. - if rule.bodyContains != "" && !bytes.Contains(body, []byte(rule.bodyContains)) { - continue - } - if rule.offset > 0 { - rule.offset-- - return nil - } - if rule.times <= 0 { - delete(fr.rules, key) - return nil - } - rule.times-- - if rule.times == 0 { - delete(fr.rules, key) - } - result := *rule - return &result - } - - return nil -} +import "github.com/databricks/cli/libs/fault" // faultEndpointHandler handles POST /__testserver/fault. -func faultEndpointHandler(fr *FaultRules) HandlerFunc { +func faultEndpointHandler(fr *fault.Rules) HandlerFunc { return func(req Request) any { - var body struct { - Pattern string `json:"pattern"` - StatusCode int `json:"status_code"` - Body string `json:"body"` - Offset int `json:"offset"` - Times int `json:"times"` - AfterHandler bool `json:"after_handler"` - BodyContains string `json:"body_contains"` - } - if err := json.Unmarshal(req.Body, &body); err != nil { + if err := fr.Register(req.Token, req.Body); err != nil { return Response{StatusCode: 400, Body: map[string]string{"error": err.Error()}} } - if body.AfterHandler { - fr.SetAfterHandler(req.Token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } else { - fr.Set(req.Token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } return Response{StatusCode: 200} } } diff --git a/libs/testserver/fault_test.go b/libs/testserver/fault_test.go index ae27b954f93..f541e3b476a 100644 --- a/libs/testserver/fault_test.go +++ b/libs/testserver/fault_test.go @@ -12,66 +12,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestFaultRulesNoMatch(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "GET /foo", 504, "body", 0, 1, "") - - assert.Nil(t, fr.Check("POST", "/foo", "tok", nil)) - assert.Nil(t, fr.Check("GET", "/bar", "tok", nil)) - assert.Nil(t, fr.Check("GET", "/foo", "other", nil)) -} - -func TestFaultRulesExactMatch(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "PUT /api/2.0/jobs/123", 504, "body", 0, 1, "") - - rule := fr.Check("PUT", "/api/2.0/jobs/123", "tok", nil) - require.NotNil(t, rule) - assert.Equal(t, 504, rule.StatusCode) - assert.Equal(t, "body", rule.Body) -} - -func TestFaultRulesWildcardMatch(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "PUT /api/2.0/permissions/pipelines/*", 504, "body", 0, 2, "") - - assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/abc", "tok", nil)) - assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) - assert.Nil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) // exhausted -} - -func TestFaultRulesOffset(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "GET /foo", 504, "body", 2, 1, "") - - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 2→1 - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 1→0 - assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) // fires - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted -} - -func TestFaultRulesTimes(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "GET /foo", 504, "body", 0, 3, "") - - for range 3 { - assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) - } - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted -} - -func TestFaultRulesBodyContains(t *testing.T) { - fr := testserver.NewFaultRules() - fr.Set("tok", "POST /api/2.0/workspace/import", 403, "body", 0, 1, "state/resources.json") - - // Body without the substring must not match and must not consume the budget. - assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/deploy.lock"))) - // Body with the substring fires. - assert.NotNil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) - // Budget was only consumed by the matching request. - assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) -} - func TestServerFaultAfterHandlerKeepsTheHandlersEffect(t *testing.T) { var calls atomic.Int32 server := testserver.New(t) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 5170b318818..a098c1de1c5 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -18,6 +18,7 @@ import ( "sync" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/fault" "github.com/databricks/cli/libs/testserver/testsql" ) @@ -67,7 +68,7 @@ type Server struct { mu sync.Mutex kills *killRules - faults *FaultRules + faults *fault.Rules sqlHandler *testsql.Handler @@ -106,6 +107,7 @@ func NewRequest(t testutil.TestingT, r *http.Request, fakeWorkspace *FakeWorkspa return Request{ Method: r.Method, + Token: getToken(r), URL: r.URL, Headers: r.Header, Body: body, @@ -227,7 +229,7 @@ func getHeaders(value []byte) http.Header { func New(t testutil.TestingT) *Server { router := NewRouter() kills := newKillRules() - faults := NewFaultRules() + faults := fault.NewRules() // Wrap the router so kill rules fire for ALL requests, including those with // no registered handler that would otherwise bypass serve() entirely. @@ -364,7 +366,6 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, handler HandlerFu request := NewRequest(s.t, r, fakeWorkspace) request.Vars = vars - request.Token = token if s.RequestCallback != nil { s.RequestCallback(&request) From 079f9f3998b8174e91c6564c6046fe831c12d9ed Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 28 Oct 2026 16:19:16 +0000 Subject: [PATCH 2/3] Add fault injection to cloud acceptance test proxy --- acceptance/README.md | 13 + acceptance/acceptance_test.go | 2 + acceptance/bin/fault.py | 3 +- acceptance/internal/config.go | 10 +- acceptance/internal/prepare_server.go | 13 +- acceptance/selftest/fault_proxy/out.test.toml | 1 + acceptance/selftest/fault_proxy/output.txt | 38 +++ acceptance/selftest/fault_proxy/script | 22 ++ acceptance/selftest/fault_proxy/test.toml | 4 + libs/testproxy/server.go | 265 +++++++++--------- libs/testproxy/server_test.go | 210 ++++++++++++++ 11 files changed, 430 insertions(+), 151 deletions(-) create mode 100644 acceptance/selftest/fault_proxy/out.test.toml create mode 100644 acceptance/selftest/fault_proxy/output.txt create mode 100644 acceptance/selftest/fault_proxy/script create mode 100644 acceptance/selftest/fault_proxy/test.toml create mode 100644 libs/testproxy/server_test.go diff --git a/acceptance/README.md b/acceptance/README.md index 9cd4d2ed0a0..9c5ce55403e 100644 --- a/acceptance/README.md +++ b/acceptance/README.md @@ -15,6 +15,19 @@ The cloud run is keyed off `CLOUD_ENV`; see `getSkipReason` in `acceptance/accep To run tests against a real workspace: `deco env run -i -n aws-prod-ucws -- ` (requires the `deco` tool and access to a test env). +## Fault injection + +Set `Proxy = true` in `test.toml` to use `fault.py` against either the fake server or a real workspace. Add `Cloud = true` for cloud runs. Recording is optional. + +```bash +fault.py "POST /api/2.0/workspace/mkdirs" 403 0 1 +fault.py --after-handler "POST /api/2.2/jobs/run-now" 503 0 1 +``` + +A normal fault prevents the upstream call. `--after-handler` forwards the call first and replaces its response, preserving its effects. Rules match the incoming test token; the proxy uses separate upstream credentials. Offsets and counts include client retries. Use precise patterns or `--body-contains` to select a resource, and clean up real resources after testing a failed response. + +The proxy requires localhost access and cannot be used with `RunsOnDbr`. + ## Authoring To author a test, diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 29f1ccde3a3..a333e17a322 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -919,6 +919,8 @@ func runTest(t *testing.T, args := []string{"bash", "-euo", "pipefail", EntryPointScript} cmd := exec.CommandContext(ctx, args[0], args[1:]...) + require.False(t, isTruePtr(config.RunsOnDbr) && LogRequests, + "RunsOnDbr cannot be combined with -logrequests: the proxy requires localhost access") cfg, user, workspaceID := internal.PrepareServerAndClient(t, config, LogRequests, outputDir, testEnv) testdiff.PrepareReplacementsUser(t, &repls, user) testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) diff --git a/acceptance/bin/fault.py b/acceptance/bin/fault.py index 5b75c0a5907..6ea3c748ed0 100755 --- a/acceptance/bin/fault.py +++ b/acceptance/bin/fault.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Set up a fault rule on the testserver for the current test token. +"""Set up a fault rule on the testserver or test proxy for the current test token. Usage: fault.py [--after-handler] [--body-contains SUBSTR] PATTERN STATUS_CODE OFFSET TIMES [ERROR_CODE] @@ -16,6 +16,7 @@ The rule is scoped to the current DATABRICKS_TOKEN so it only affects the test that registers it, even when tests share a server. +Cloud runs require Proxy = true in test.toml. """ import json diff --git a/acceptance/internal/config.go b/acceptance/internal/config.go index 6ea79a9578a..8cbe0f68b42 100644 --- a/acceptance/internal/config.go +++ b/acceptance/internal/config.go @@ -78,8 +78,8 @@ type TestConfig struct { // out.requests.txt RecordRequests *bool - // If true, local runs route the CLI through the recording proxy (libs/testproxy) - // instead of straight to the testserver, matching the cloud topology. + // If true, route the CLI through libs/testproxy for fault injection, + // independently of recording. Supported locally and on cloud. Proxy *bool // List of request headers to include when recording requests. @@ -252,10 +252,10 @@ func LoadConfig(t *testing.T, dir string) (TestConfig, string) { // validateConfig checks for incompatible configuration combinations. func validateConfig(t *testing.T, config TestConfig, configPath string) { - // RunsOnDbr and RecordRequests are incompatible because serverless does not + // RunsOnDbr cannot use the proxy because serverless does not // allow access to localhost ports, which the test proxy server requires. - if isTruePtr(config.RunsOnDbr) && isTruePtr(config.RecordRequests) { - t.Fatalf("Invalid config %s: RunsOnDbr and RecordRequests cannot both be true. "+ + if isTruePtr(config.RunsOnDbr) && (isTruePtr(config.RecordRequests) || isTruePtr(config.Proxy)) { + t.Fatalf("Invalid config %s: RunsOnDbr cannot be combined with RecordRequests or Proxy. "+ "Serverless does not allow access to localhost ports, which the test proxy server requires.", configPath) } diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 3890aa52d5a..59b564372de 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -129,10 +129,10 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o cfg := w.Config // If we are running in a cloud environment AND we need to intercept requests - // (for recording or logging), start a proxy server. - if recordRequests || logRequests { - // An empty config resolves the real workspace from the environment. - host := startProxyServer(t, &sdkconfig.Config{}, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + // (for fault injection, recording or logging), start a proxy server. + if isTruePtr(config.Proxy) || recordRequests || logRequests { + // Reuse the resolved cloud config for upstream auth and workspace routing. + host := startProxyServer(t, cfg, token, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) cfg = &sdkconfig.Config{ Host: host, Token: token, @@ -155,7 +155,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o host := startProxyServer(t, &sdkconfig.Config{ Host: upstream, Token: token, - }, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + }, token, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) cfg := &sdkconfig.Config{ Host: host, @@ -284,12 +284,13 @@ func startLocalServer(t *testing.T, func startProxyServer(t *testing.T, upstream *sdkconfig.Config, + token string, recordRequests bool, logRequests bool, includeHeaders []string, outputDir string, ) string { - s := testproxy.New(t, upstream) + s := testproxy.New(t, upstream, token) // Record API requests in out.requests.txt if RecordRequests is true in test.toml. if recordRequests { diff --git a/acceptance/selftest/fault_proxy/out.test.toml b/acceptance/selftest/fault_proxy/out.test.toml new file mode 100644 index 00000000000..c7c6f58ed6e --- /dev/null +++ b/acceptance/selftest/fault_proxy/out.test.toml @@ -0,0 +1 @@ +Cloud = true diff --git a/acceptance/selftest/fault_proxy/output.txt b/acceptance/selftest/fault_proxy/output.txt new file mode 100644 index 00000000000..4ae843483cc --- /dev/null +++ b/acceptance/selftest/fault_proxy/output.txt @@ -0,0 +1,38 @@ + +=== a fault before forwarding prevents the mutation +>>> musterr [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/blocked +Error: Fault injected by test. + +Host: [DATABRICKS_URL] +Auth type: Personal Access Token (pat) + +Next steps: + - Verify you have the required permissions for this operation + - Check your identity: databricks auth describe + - Consider setting up a profile: databricks auth login --profile +blocked directory does not exist + +=== an after-handler fault preserves the upstream mutation +>>> musterr [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/committed +Error: Fault injected by test. + +Host: [DATABRICKS_URL] +Auth type: Personal Access Token (pat) + +Next steps: + - Verify you have the required permissions for this operation + - Check your identity: databricks auth describe + - Consider setting up a profile: databricks auth login --profile + +>>> [CLI] workspace get-status /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/committed +{ + "object_type": "DIRECTORY" +} + +=== the CLI recovers from a retryable error after the mutation +>>> [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/retried + +>>> [CLI] workspace get-status /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/retried +{ + "object_type": "DIRECTORY" +} diff --git a/acceptance/selftest/fault_proxy/script b/acceptance/selftest/fault_proxy/script new file mode 100644 index 00000000000..81b0aaef615 --- /dev/null +++ b/acceptance/selftest/fault_proxy/script @@ -0,0 +1,22 @@ +root="/Users/$CURRENT_USER_NAME/fault-proxy-$UNIQUE_NAME" +cleanup() { + $CLI workspace delete "$root" --recursive +} +$CLI workspace mkdirs "$root" +trap cleanup EXIT + +title "a fault before forwarding prevents the mutation" +fault.py --body-contains "$root/blocked" "POST /api/2.0/workspace/mkdirs" 403 0 1 +trace musterr $CLI workspace mkdirs "$root/blocked" +musterr $CLI workspace get-status "$root/blocked" >/dev/null 2>&1 +echo "blocked directory does not exist" + +title "an after-handler fault preserves the upstream mutation" +fault.py --after-handler --body-contains "$root/committed" "POST /api/2.0/workspace/mkdirs" 403 0 1 +trace musterr $CLI workspace mkdirs "$root/committed" +trace $CLI workspace get-status "$root/committed" | jq '{object_type}' + +title "the CLI recovers from a retryable error after the mutation" +fault.py --after-handler --body-contains "$root/retried" "POST /api/2.0/workspace/mkdirs" 503 0 1 +trace $CLI workspace mkdirs "$root/retried" +trace $CLI workspace get-status "$root/retried" | jq '{object_type}' diff --git a/acceptance/selftest/fault_proxy/test.toml b/acceptance/selftest/fault_proxy/test.toml new file mode 100644 index 00000000000..dd645d54e45 --- /dev/null +++ b/acceptance/selftest/fault_proxy/test.toml @@ -0,0 +1,4 @@ +Cloud = true +Proxy = true +RecordRequests = false +Env.MSYS_NO_PATHCONV = '1' diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index be5bcfd4aa7..bb444268b96 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -2,186 +2,173 @@ package testproxy import ( "bytes" - "errors" + "crypto/tls" + "fmt" + "io" + "maps" "net/http" "net/http/httptest" + "net/http/httputil" + "net/url" + "strings" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/fault" "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/httpclient" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +const faultEndpoint = "/__testserver/fault" + type ProxyServer struct { *httptest.Server - t testutil.TestingT - - apiClient *httpclient.ApiClient + t testutil.TestingT + upstream *url.URL + transport http.RoundTripper + token string + faults *fault.Rules RequestCallback func(request *testserver.Request) ResponseCallback func(request *testserver.Request, response *testserver.EncodedResponse) } -// New creates a reverse proxy server that sits in front of the upstream -// workspace. This is useful for recording API requests and responses in -// integration tests. -// -// Note: We cannot directly proxy the request from a localhost URL to a real -// workspace as is. This is because auth resolution in the Databricks SDK relies -// what the URL actually looks like to determine the auth method to use. -// For example, in OAuth flows, the SDK can make requests to different Microsoft -// OAuth endpoints based on the nature of the URL. -// For reference, see: +// New creates a reverse proxy with token-scoped fault injection and optional recording. +// The upstream config resolves real workspace credentials; token identifies the test +// locally and is never forwarded. Resolving auth against localhost would break OAuth: // https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 -// -// An upstream with no fields set resolves a real workspace and its auth from the -// environment; an explicit host and token point at a testserver instead. -func New(t testutil.TestingT, upstream *config.Config) *ProxyServer { - s := &ProxyServer{ - t: t, - } - - clientCfg, err := config.HTTPClientConfigFromConfig(upstream) +func New(t testutil.TestingT, upstream *config.Config, token string) *ProxyServer { + require.NoError(t, upstream.EnsureResolved()) + upstreamURL, err := url.Parse(upstream.Host) require.NoError(t, err) - s.apiClient = httpclient.NewApiClient(clientCfg) - - // Set up the proxy handler as the default handler for all requests. - server := httptest.NewServer(http.HandlerFunc(s.proxyToUpstream)) - t.Cleanup(server.Close) - s.Server = server + transport := upstream.HTTPTransport + if transport == nil { + tr := http.DefaultTransport.(*http.Transport).Clone() + tr.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, InsecureSkipVerify: upstream.InsecureSkipVerify} //nolint:gosec // Honor the upstream SDK config. + t.Cleanup(tr.CloseIdleConnections) + transport = tr + } + s := &ProxyServer{ + t: t, + upstream: upstreamURL, + transport: &authenticatedTransport{upstream: upstream, next: transport}, + token: token, + faults: fault.NewRules(), + } + s.Server = httptest.NewServer(http.HandlerFunc(s.proxyToUpstream)) + t.Cleanup(s.Close) return s } -func (s *ProxyServer) reqBody(r testserver.Request) any { - // The SDK expects the query parameters to be specified in the "request body" - // argument for GET, DELETE, and HEAD requests in the .Do method. - if r.Method == http.MethodGet || r.Method == http.MethodDelete || r.Method == http.MethodHead { - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] +type authenticatedTransport struct { + upstream *config.Config + next http.RoundTripper +} + +func (tr *authenticatedTransport) RoundTrip(r *http.Request) (*http.Response, error) { + r.Header.Del("Authorization") + if err := tr.upstream.Authenticate(r); err != nil { + return nil, fmt.Errorf("authenticate upstream request: %w", err) + } + for key, value := range auth.WorkspaceIDHeaders(tr.upstream) { + r.Header.Set(key, value) + } + if tr.upstream.Headers != nil { + if err := tr.upstream.Headers(r); err != nil { + return nil, err } - return queryParams } - - // Otherwise, return the raw body. - return r.Body + return tr.next.RoundTrip(r) } func (s *ProxyServer) proxyToUpstream(w http.ResponseWriter, r *http.Request) { request := testserver.NewRequest(s.t, r, nil) - if s.RequestCallback != nil { - s.RequestCallback(&request) - } - - headers := make(map[string]string) - for k, v := range r.Header { - // Authorization headers will be set by the SDK. Do not pass them along here. - if k == "Authorization" { - continue + if strings.HasPrefix(r.URL.Path, "/__testserver/") { + if r.URL.Path != faultEndpoint || r.Method != http.MethodPost { + http.NotFound(w, r) + return } - // The default HTTP client in Go sets the Accept-Encoding header to - // "gzip". Since it is originally meant to be read by the server and - // will be set again when the SDK makes the request to the workspace, do - // not pass it along here. - if k == "Accept-Encoding" { - continue + if request.Token == "" || request.Token != s.token { + http.Error(w, "invalid test token", http.StatusUnauthorized) + return } - headers[k] = v[0] - } - - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] + if err := s.faults.Register(request.Token, request.Body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusOK) + return } - reqBody := s.reqBody(request) - respBody := bytes.Buffer{} - - // List of response headers to pass through to the client. - // We can extend this code to pass all X-Databricks headers if we introduce a - // more powerful visitor like `WithRequestHeaders` to the SDK. That is left - // as a future exercise for now. - includeResponseHeaders := []string{"X-Databricks-Org-Id"} - responseHeaders := make(map[string]*string) - - visitors := []httpclient.DoOption{ - httpclient.WithRequestHeaders(headers), - httpclient.WithQueryParameters(queryParams), - httpclient.WithRequestData(reqBody), - httpclient.WithResponseUnmarshal(&respBody), + if s.RequestCallback != nil { + s.RequestCallback(&request) } - - for _, header := range includeResponseHeaders { - responseHeaders[header] = new(string) - visitors = append(visitors, httpclient.WithResponseHeader(header, responseHeaders[header])) + rule := s.faults.Check(r.Method, r.URL.Path, request.Token, request.Body) + if rule != nil && !rule.AfterHandler { + s.writeResponse(w, &request, injectedResponse(rule)) + return } - err := s.apiClient.Do(s.t.Context(), r.Method, r.URL.Path, - visitors..., - ) - - var encodedResponse *testserver.EncodedResponse - - // API errors from the SDK are of type [apierr.APIError]. Forward the raw - // response bytes verbatim — including any error details — so callers see - // exactly what the workspace returned. Re-marshalling from the parsed - // APIError would drop fields the SDK doesn't surface (e.g. metadata in - // details[]) and silently break callers that inspect them. - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { - rw := apiErr.ResponseWrapper - if rw == nil { - // The SDK populates ResponseWrapper for every APIError produced - // from a real HTTP response. If this ever fires the SDK changed - // shape and we need to revisit how we forward error bodies. - panic("apierr.APIError has no ResponseWrapper") - } - encodedResponse = &testserver.EncodedResponse{ - StatusCode: rw.Response.StatusCode, - Body: rw.DebugBytes, - } - // Visitors registered via WithResponseHeader are not invoked when - // the SDK returns an error, so populate the include list directly - // from the original response headers. - for _, header := range includeResponseHeaders { - if v := rw.Response.Header.Get(header); v != "" { - *responseHeaders[header] = v + r.Body = io.NopCloser(bytes.NewReader(request.Body)) + proxy := httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(s.upstream) + // Let the transport negotiate compression so response logging sees decoded bytes. + pr.Out.Header.Del("Accept-Encoding") + // Preserve the client's query encoding and repeated parameters. + pr.Out.URL.RawQuery = pr.In.URL.RawQuery + }, + Transport: s.transport, + ModifyResponse: func(response *http.Response) error { + body, err := io.ReadAll(response.Body) + response.Body.Close() + if err != nil { + return err } - } - } - - // Something else went wrong. - if encodedResponse == nil && err != nil { - encodedResponse = &testserver.EncodedResponse{ - StatusCode: 500, - Body: []byte(err.Error()), - } - } - - // Successful response - if encodedResponse == nil { - encodedResponse = &testserver.EncodedResponse{ - StatusCode: 200, - Body: respBody.Bytes(), - } + encoded := testserver.EncodedResponse{ + StatusCode: response.StatusCode, + Headers: response.Header, + Body: body, + } + if rule != nil { + encoded = injectedResponse(rule) + } + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &encoded) + } + response.StatusCode = encoded.StatusCode + response.Header = encoded.Headers + response.Body = io.NopCloser(bytes.NewReader(encoded.Body)) + response.ContentLength = int64(len(encoded.Body)) + return nil + }, + ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) { + s.writeResponse(w, &request, testserver.EncodedResponse{ + StatusCode: http.StatusBadGateway, + Body: []byte(err.Error()), + }) + }, } + proxy.ServeHTTP(w, r) +} - // Send response to client. - if len(responseHeaders) > 0 { - for k, v := range responseHeaders { - w.Header()[k] = []string{*v} - } +func injectedResponse(rule *fault.Rule) testserver.EncodedResponse { + return testserver.EncodedResponse{ + StatusCode: rule.StatusCode, + Headers: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(rule.Body), } +} - w.WriteHeader(encodedResponse.StatusCode) - - _, err = w.Write(encodedResponse.Body) - assert.NoError(s.t, err) - +func (s *ProxyServer) writeResponse(w http.ResponseWriter, request *testserver.Request, response testserver.EncodedResponse) { if s.ResponseCallback != nil { - s.ResponseCallback(&request, encodedResponse) + s.ResponseCallback(request, &response) } + maps.Copy(w.Header(), response.Headers) + w.WriteHeader(response.StatusCode) + _, err := w.Write(response.Body) + assert.NoError(s.t, err) } diff --git a/libs/testproxy/server_test.go b/libs/testproxy/server_test.go new file mode 100644 index 00000000000..eff5c16b6a3 --- /dev/null +++ b/libs/testproxy/server_test.go @@ -0,0 +1,210 @@ +package testproxy_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/testproxy" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testToken = "test-token" + +func newProxy(t *testing.T, handler http.HandlerFunc) *testproxy.ProxyServer { + t.Helper() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"workspace_id":"123"}`)) + return + } + handler(w, r) + })) + t.Cleanup(upstream.Close) + return testproxy.New(t, &config.Config{ + Host: upstream.URL, Token: "upstream-token", AuthType: "pat", ConfigFile: os.DevNull, + WorkspaceID: "123", + }, testToken) +} + +func send(t *testing.T, proxy *testproxy.ProxyServer, method, path, token string, body []byte) (int, http.Header, []byte) { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), method, proxy.URL+path, bytes.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + response, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer response.Body.Close() + data, err := io.ReadAll(response.Body) + require.NoError(t, err) + return response.StatusCode, response.Header, data +} + +func register(t *testing.T, proxy *testproxy.ProxyServer, after bool) { + t.Helper() + data, err := json.Marshal(map[string]any{ + "pattern": "POST /api/write", "status_code": 503, "body": `{"error_code":"INJECTED"}`, + "offset": 1, "times": 1, "body_contains": "target", "after_handler": after, + }) + require.NoError(t, err) + status, _, _ := send(t, proxy, "POST", "/__testserver/fault", testToken, data) + require.Equal(t, 200, status) +} + +func TestProxyFaults(t *testing.T) { + for _, after := range []bool{false, true} { + t.Run(map[bool]string{false: "before", true: "after"}[after], func(t *testing.T) { + var calls, recorded, responses atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/write" { + w.WriteHeader(http.StatusNotFound) + return + } + assert.Equal(t, "Bearer upstream-token", r.Header.Get("Authorization")) + assert.Equal(t, "123", r.Header.Get("X-Databricks-Workspace-Id")) + calls.Add(1) + w.Header().Set("Retry-After", "23") + w.Header().Set("Content-Type", "text/plain") + w.Header().Set("Content-Length", "9") + w.WriteHeader(http.StatusCreated) + _, _ = w.Write([]byte("committed")) + }) + proxy.RequestCallback = func(*testserver.Request) { recorded.Add(1) } + proxy.ResponseCallback = func(*testserver.Request, *testserver.EncodedResponse) { responses.Add(1) } + register(t, proxy, after) + assert.Zero(t, calls.Load()) + assert.Zero(t, recorded.Load()) + + // Neither another token nor another body consumes the matching budget. + for _, request := range []struct{ token, body string }{{"other-token", "target"}, {testToken, "other"}, {testToken, "target"}} { + status, _, body := send(t, proxy, "POST", "/api/write", request.token, []byte(request.body)) + assert.Equal(t, 201, status) + assert.Equal(t, "committed", string(body)) + } + status, headers, body := send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 503, status) + assert.JSONEq(t, `{"error_code":"INJECTED"}`, string(body)) + assert.Equal(t, "application/json", headers.Get("Content-Type")) + assert.Empty(t, headers.Get("Retry-After")) + if after { + assert.Equal(t, int32(4), calls.Load()) + } else { + assert.Equal(t, int32(3), calls.Load()) + } + status, _, _ = send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 201, status) + assert.Equal(t, int32(5), recorded.Load()) + assert.Equal(t, int32(5), responses.Load()) + }) + } +} + +func TestProxyControlDoesNotReachUpstream(t *testing.T) { + var calls atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, _ *http.Request) { calls.Add(1); w.WriteHeader(http.StatusNotFound) }) + for _, tc := range []struct { + method, path, token string + status int + }{ + {"POST", "/__testserver/fault", "other", 401}, + {"POST", "/__testserver/fault", "", 401}, + {"POST", "/__testserver/kill", testToken, 404}, + {"GET", "/__testserver/fault", testToken, 404}, + {"POST", "/__testserver/fault", testToken, 400}, + } { + t.Run(tc.method+tc.path+tc.token, func(t *testing.T) { + status, _, _ := send(t, proxy, tc.method, tc.path, tc.token, []byte("invalid JSON")) + assert.Equal(t, tc.status, status) + }) + } + assert.Zero(t, calls.Load()) +} + +func TestProxyPreservesUpstreamResponse(t *testing.T) { + var calls atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/write" { + w.WriteHeader(http.StatusNotFound) + return + } + calls.Add(1) + assert.Equal(t, "x=a%2Fb&x=c+d", r.URL.RawQuery) + body, err := io.ReadAll(r.Body) + assert.NoError(t, err) + assert.Equal(t, []byte("multipart\x00bytes\r\n"), body) + w.Header().Set("Retry-After", "7") + w.Header().Add("X-Example", "one") + w.Header().Add("X-Example", "two") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"details":[{"reason":"preserved"}]}`)) + }) + status, headers, body := send(t, proxy, "POST", "/api/write?x=a%2Fb&x=c+d", testToken, []byte("multipart\x00bytes\r\n")) + assert.Equal(t, 429, status) + assert.Equal(t, "7", headers.Get("Retry-After")) + assert.Equal(t, []string{"one", "two"}, headers.Values("X-Example")) + assert.JSONEq(t, `{"details":[{"reason":"preserved"}]}`, string(body)) + assert.Equal(t, int32(1), calls.Load()) +} + +func TestProxyDoesNotMaskTransportFailure(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"workspace_id":"123"}`)) + })) + t.Cleanup(upstream.Close) + proxy := testproxy.New(t, &config.Config{ + Host: upstream.URL, Token: "upstream-token", AuthType: "pat", ConfigFile: os.DevNull, + }, testToken) + upstream.Close() + data := []byte(`{"pattern":"POST /api/write","status_code":403,"body":"injected","after_handler":true,"times":1}`) + status, _, _ := send(t, proxy, "POST", "/__testserver/fault", testToken, data) + require.Equal(t, 200, status) + status, _, body := send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 502, status) + assert.NotEqual(t, "injected", string(body)) +} + +func TestProxyCancellation(t *testing.T) { + started := make(chan struct{}) + cancelled := make(chan struct{}) + proxy := newProxy(t, func(_ http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + close(cancelled) + }) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxy.URL+"/api/wait", nil) + require.NoError(t, err) + done := make(chan error, 1) + go func() { + response, err := http.DefaultClient.Do(req) + if response != nil { + response.Body.Close() + } + done <- err + }() + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("upstream request did not start") + } + cancel() + require.ErrorIs(t, <-done, context.Canceled) + select { + case <-cancelled: + case <-time.After(5 * time.Second): + t.Fatal("upstream request was not cancelled") + } +} From f073dd4f838808e7020f5c5a60962d95383934a3 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 28 Oct 2026 16:54:37 +0000 Subject: [PATCH 3/3] Keep shared fault rules and tests in testserver --- libs/fault/fault.go | 132 --------------------------------- libs/fault/fault_test.go | 69 ------------------ libs/testproxy/server.go | 7 +- libs/testserver/fault.go | 133 +++++++++++++++++++++++++++++++++- libs/testserver/fault_test.go | 60 +++++++++++++++ libs/testserver/server.go | 6 +- 6 files changed, 197 insertions(+), 210 deletions(-) delete mode 100644 libs/fault/fault.go delete mode 100644 libs/fault/fault_test.go diff --git a/libs/fault/fault.go b/libs/fault/fault.go deleted file mode 100644 index 9da0d579d72..00000000000 --- a/libs/fault/fault.go +++ /dev/null @@ -1,132 +0,0 @@ -package fault - -import ( - "bytes" - "encoding/json" - "strings" - "sync" -) - -type faultRuleKey struct { - token string - pattern string -} - -// Rule describes a single injected fault: HTTP status, body, and remaining fire count. -type Rule struct { - StatusCode int - Body string - // AfterHandler keeps the handler's state change and replaces only its response. - AfterHandler bool - offset int - times int - - // bodyContains, if non-empty, additionally requires the request body to - // contain this substring for the rule to fire. /workspace/import routes - // every upload through the same method+path, so the only way to target a - // single file's upload is by matching its multipart "path" form field. - bodyContains string -} - -// Rules holds the active fault injection rules for a test server. -type Rules struct { - mu sync.Mutex - rules map[faultRuleKey]*Rule -} - -// NewRules returns an empty Rules. -func NewRules() *Rules { - return &Rules{rules: make(map[faultRuleKey]*Rule)} -} - -// Set registers or replaces a fault rule for the given token and pattern. -func (fr *Rules) Set(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { - fr.set(token, pattern, statusCode, body, offset, times, false, bodyContains) -} - -// SetAfterHandler is like Set, but the handler runs first so its effect is kept. -func (fr *Rules) SetAfterHandler(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { - fr.set(token, pattern, statusCode, body, offset, times, true, bodyContains) -} - -func (fr *Rules) set(token, pattern string, statusCode int, body string, offset, times int, afterHandler bool, bodyContains string) { - fr.mu.Lock() - defer fr.mu.Unlock() - fr.rules[faultRuleKey{token: token, pattern: pattern}] = &Rule{ - StatusCode: statusCode, - Body: body, - AfterHandler: afterHandler, - offset: offset, - times: times, - bodyContains: bodyContains, - } -} - -// Check returns a matching fault rule and advances its counters, or nil if no rule matches. -// Pattern supports a trailing "*" wildcard, e.g. "PUT /api/2.0/permissions/pipelines/*". -// A rule with a non-empty bodyContains only matches when body contains that substring. -func (fr *Rules) Check(method, path, token string, body []byte) *Rule { - requestPattern := method + " " + path - - fr.mu.Lock() - defer fr.mu.Unlock() - - for key, rule := range fr.rules { - if key.token != token { - continue - } - rulePattern := key.pattern - var matched bool - if strings.HasSuffix(rulePattern, "*") { - matched = strings.HasPrefix(requestPattern, rulePattern[:len(rulePattern)-1]) - } else { - matched = requestPattern == rulePattern - } - if !matched { - continue - } - // A non-matching body must not consume the rule's offset/times budget, - // so this check precedes the counter updates below. - if rule.bodyContains != "" && !bytes.Contains(body, []byte(rule.bodyContains)) { - continue - } - if rule.offset > 0 { - rule.offset-- - return nil - } - if rule.times <= 0 { - delete(fr.rules, key) - return nil - } - rule.times-- - if rule.times == 0 { - delete(fr.rules, key) - } - result := *rule - return &result - } - - return nil -} - -// Register decodes and registers a rule for token using the fault.py wire format. -func (fr *Rules) Register(token string, data []byte) error { - var body struct { - Pattern string `json:"pattern"` - StatusCode int `json:"status_code"` - Body string `json:"body"` - Offset int `json:"offset"` - Times int `json:"times"` - AfterHandler bool `json:"after_handler"` - BodyContains string `json:"body_contains"` - } - if err := json.Unmarshal(data, &body); err != nil { - return err - } - if body.AfterHandler { - fr.SetAfterHandler(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } else { - fr.Set(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } - return nil -} diff --git a/libs/fault/fault_test.go b/libs/fault/fault_test.go deleted file mode 100644 index 202f7761a08..00000000000 --- a/libs/fault/fault_test.go +++ /dev/null @@ -1,69 +0,0 @@ -package fault_test - -import ( - "testing" - - "github.com/databricks/cli/libs/fault" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestFaultRulesNoMatch(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "GET /foo", 504, "body", 0, 1, "") - - assert.Nil(t, fr.Check("POST", "/foo", "tok", nil)) - assert.Nil(t, fr.Check("GET", "/bar", "tok", nil)) - assert.Nil(t, fr.Check("GET", "/foo", "other", nil)) -} - -func TestFaultRulesExactMatch(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "PUT /api/2.0/jobs/123", 504, "body", 0, 1, "") - - rule := fr.Check("PUT", "/api/2.0/jobs/123", "tok", nil) - require.NotNil(t, rule) - assert.Equal(t, 504, rule.StatusCode) - assert.Equal(t, "body", rule.Body) -} - -func TestFaultRulesWildcardMatch(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "PUT /api/2.0/permissions/pipelines/*", 504, "body", 0, 2, "") - - assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/abc", "tok", nil)) - assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) - assert.Nil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) // exhausted -} - -func TestFaultRulesOffset(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "GET /foo", 504, "body", 2, 1, "") - - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 2→1 - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 1→0 - assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) // fires - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted -} - -func TestFaultRulesTimes(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "GET /foo", 504, "body", 0, 3, "") - - for range 3 { - assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) - } - assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted -} - -func TestFaultRulesBodyContains(t *testing.T) { - fr := fault.NewRules() - fr.Set("tok", "POST /api/2.0/workspace/import", 403, "body", 0, 1, "state/resources.json") - - // Body without the substring must not match and must not consume the budget. - assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/deploy.lock"))) - // Body with the substring fires. - assert.NotNil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) - // Budget was only consumed by the matching request. - assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) -} diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index bb444268b96..6cb547c6f32 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -14,7 +14,6 @@ import ( "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/auth" - "github.com/databricks/cli/libs/fault" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go/config" "github.com/stretchr/testify/assert" @@ -30,7 +29,7 @@ type ProxyServer struct { upstream *url.URL transport http.RoundTripper token string - faults *fault.Rules + faults *testserver.FaultRules RequestCallback func(request *testserver.Request) ResponseCallback func(request *testserver.Request, response *testserver.EncodedResponse) } @@ -56,7 +55,7 @@ func New(t testutil.TestingT, upstream *config.Config, token string) *ProxyServe upstream: upstreamURL, transport: &authenticatedTransport{upstream: upstream, next: transport}, token: token, - faults: fault.NewRules(), + faults: testserver.NewFaultRules(), } s.Server = httptest.NewServer(http.HandlerFunc(s.proxyToUpstream)) t.Cleanup(s.Close) @@ -155,7 +154,7 @@ func (s *ProxyServer) proxyToUpstream(w http.ResponseWriter, r *http.Request) { proxy.ServeHTTP(w, r) } -func injectedResponse(rule *fault.Rule) testserver.EncodedResponse { +func injectedResponse(rule *testserver.FaultRule) testserver.EncodedResponse { return testserver.EncodedResponse{ StatusCode: rule.StatusCode, Headers: http.Header{"Content-Type": {"application/json"}}, diff --git a/libs/testserver/fault.go b/libs/testserver/fault.go index fff8db13a8a..eea6c117e79 100644 --- a/libs/testserver/fault.go +++ b/libs/testserver/fault.go @@ -1,9 +1,138 @@ package testserver -import "github.com/databricks/cli/libs/fault" +import ( + "bytes" + "encoding/json" + "strings" + "sync" +) + +type faultRuleKey struct { + token string + pattern string +} + +// FaultRule describes a single injected fault: HTTP status, body, and remaining fire count. +type FaultRule struct { + StatusCode int + Body string + // AfterHandler keeps the handler's state change and replaces only its response. + AfterHandler bool + offset int + times int + + // bodyContains, if non-empty, additionally requires the request body to + // contain this substring for the rule to fire. /workspace/import routes + // every upload through the same method+path, so the only way to target a + // single file's upload is by matching its multipart "path" form field. + bodyContains string +} + +// FaultRules holds the active fault injection rules for a test server. +type FaultRules struct { + mu sync.Mutex + rules map[faultRuleKey]*FaultRule +} + +// NewFaultRules returns an empty FaultRules. +func NewFaultRules() *FaultRules { + return &FaultRules{rules: make(map[faultRuleKey]*FaultRule)} +} + +// Set registers or replaces a fault rule for the given token and pattern. +func (fr *FaultRules) Set(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { + fr.set(token, pattern, statusCode, body, offset, times, false, bodyContains) +} + +// SetAfterHandler is like Set, but the handler runs first so its effect is kept. +func (fr *FaultRules) SetAfterHandler(token, pattern string, statusCode int, body string, offset, times int, bodyContains string) { + fr.set(token, pattern, statusCode, body, offset, times, true, bodyContains) +} + +func (fr *FaultRules) set(token, pattern string, statusCode int, body string, offset, times int, afterHandler bool, bodyContains string) { + fr.mu.Lock() + defer fr.mu.Unlock() + fr.rules[faultRuleKey{token: token, pattern: pattern}] = &FaultRule{ + StatusCode: statusCode, + Body: body, + AfterHandler: afterHandler, + offset: offset, + times: times, + bodyContains: bodyContains, + } +} + +// Check returns a matching fault rule and advances its counters, or nil if no rule matches. +// Pattern supports a trailing "*" wildcard, e.g. "PUT /api/2.0/permissions/pipelines/*". +// A rule with a non-empty bodyContains only matches when body contains that substring. +func (fr *FaultRules) Check(method, path, token string, body []byte) *FaultRule { + requestPattern := method + " " + path + + fr.mu.Lock() + defer fr.mu.Unlock() + + for key, rule := range fr.rules { + if key.token != token { + continue + } + rulePattern := key.pattern + var matched bool + if strings.HasSuffix(rulePattern, "*") { + matched = strings.HasPrefix(requestPattern, rulePattern[:len(rulePattern)-1]) + } else { + matched = requestPattern == rulePattern + } + if !matched { + continue + } + // A non-matching body must not consume the rule's offset/times budget, + // so this check precedes the counter updates below. + if rule.bodyContains != "" && !bytes.Contains(body, []byte(rule.bodyContains)) { + continue + } + if rule.offset > 0 { + rule.offset-- + return nil + } + if rule.times <= 0 { + delete(fr.rules, key) + return nil + } + rule.times-- + if rule.times == 0 { + delete(fr.rules, key) + } + result := *rule + return &result + } + + return nil +} + +// Register decodes and registers a rule for token using the fault.py wire format. +func (fr *FaultRules) Register(token string, data []byte) error { + var body struct { + Pattern string `json:"pattern"` + StatusCode int `json:"status_code"` + Body string `json:"body"` + Offset int `json:"offset"` + Times int `json:"times"` + AfterHandler bool `json:"after_handler"` + BodyContains string `json:"body_contains"` + } + if err := json.Unmarshal(data, &body); err != nil { + return err + } + if body.AfterHandler { + fr.SetAfterHandler(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } else { + fr.Set(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } + return nil +} // faultEndpointHandler handles POST /__testserver/fault. -func faultEndpointHandler(fr *fault.Rules) HandlerFunc { +func faultEndpointHandler(fr *FaultRules) HandlerFunc { return func(req Request) any { if err := fr.Register(req.Token, req.Body); err != nil { return Response{StatusCode: 400, Body: map[string]string{"error": err.Error()}} diff --git a/libs/testserver/fault_test.go b/libs/testserver/fault_test.go index f541e3b476a..ae27b954f93 100644 --- a/libs/testserver/fault_test.go +++ b/libs/testserver/fault_test.go @@ -12,6 +12,66 @@ import ( "github.com/stretchr/testify/require" ) +func TestFaultRulesNoMatch(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "GET /foo", 504, "body", 0, 1, "") + + assert.Nil(t, fr.Check("POST", "/foo", "tok", nil)) + assert.Nil(t, fr.Check("GET", "/bar", "tok", nil)) + assert.Nil(t, fr.Check("GET", "/foo", "other", nil)) +} + +func TestFaultRulesExactMatch(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "PUT /api/2.0/jobs/123", 504, "body", 0, 1, "") + + rule := fr.Check("PUT", "/api/2.0/jobs/123", "tok", nil) + require.NotNil(t, rule) + assert.Equal(t, 504, rule.StatusCode) + assert.Equal(t, "body", rule.Body) +} + +func TestFaultRulesWildcardMatch(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "PUT /api/2.0/permissions/pipelines/*", 504, "body", 0, 2, "") + + assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/abc", "tok", nil)) + assert.NotNil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) + assert.Nil(t, fr.Check("PUT", "/api/2.0/permissions/pipelines/xyz", "tok", nil)) // exhausted +} + +func TestFaultRulesOffset(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "GET /foo", 504, "body", 2, 1, "") + + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 2→1 + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // offset 1→0 + assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) // fires + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted +} + +func TestFaultRulesTimes(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "GET /foo", 504, "body", 0, 3, "") + + for range 3 { + assert.NotNil(t, fr.Check("GET", "/foo", "tok", nil)) + } + assert.Nil(t, fr.Check("GET", "/foo", "tok", nil)) // exhausted +} + +func TestFaultRulesBodyContains(t *testing.T) { + fr := testserver.NewFaultRules() + fr.Set("tok", "POST /api/2.0/workspace/import", 403, "body", 0, 1, "state/resources.json") + + // Body without the substring must not match and must not consume the budget. + assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/deploy.lock"))) + // Body with the substring fires. + assert.NotNil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) + // Budget was only consumed by the matching request. + assert.Nil(t, fr.Check("POST", "/api/2.0/workspace/import", "tok", []byte("path=state/resources.json"))) +} + func TestServerFaultAfterHandlerKeepsTheHandlersEffect(t *testing.T) { var calls atomic.Int32 server := testserver.New(t) diff --git a/libs/testserver/server.go b/libs/testserver/server.go index a098c1de1c5..bb74de25f40 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -18,7 +18,6 @@ import ( "sync" "github.com/databricks/cli/internal/testutil" - "github.com/databricks/cli/libs/fault" "github.com/databricks/cli/libs/testserver/testsql" ) @@ -68,7 +67,7 @@ type Server struct { mu sync.Mutex kills *killRules - faults *fault.Rules + faults *FaultRules sqlHandler *testsql.Handler @@ -229,7 +228,7 @@ func getHeaders(value []byte) http.Header { func New(t testutil.TestingT) *Server { router := NewRouter() kills := newKillRules() - faults := fault.NewRules() + faults := NewFaultRules() // Wrap the router so kill rules fire for ALL requests, including those with // no registered handler that would otherwise bypass serve() entirely. @@ -366,6 +365,7 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, handler HandlerFu request := NewRequest(s.t, r, fakeWorkspace) request.Vars = vars + request.Token = token if s.RequestCallback != nil { s.RequestCallback(&request)