[Fix] Support custom retry logic per method #1081
+375
−91
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,54 @@ | ||
| package config | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "io" | ||
| "net/http" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/databricks/databricks-sdk-go/httpclient" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| type hc func(r *http.Request) (*http.Response, error) | ||
|
|
||
| func (cb hc) RoundTrip(r *http.Request) (*http.Response, error) { | ||
| return cb(r) | ||
| } | ||
|
|
||
| func (cb hc) SkipRetryOnIO() bool { | ||
| return true | ||
| } | ||
|
|
||
| func TestApiClient_RetriesGetPermissionsOnGatewayTimeout(t *testing.T) { | ||
| requestCount := 0 | ||
| c := &Config{ | ||
| HTTPTransport: hc(func(r *http.Request) (*http.Response, error) { | ||
| initialRequestCount := requestCount | ||
| requestCount++ | ||
| if initialRequestCount == 0 { | ||
| return &http.Response{ | ||
| Request: r, | ||
| StatusCode: http.StatusGatewayTimeout, | ||
| Body: io.NopCloser(strings.NewReader( | ||
| fmt.Sprintf(`{"error_code":"TEMPORARILY_UNAVAILABLE", "message":"The service at %s is taking too long to process your request. Please try again later or try a faster operation."}`, r.URL))), | ||
| }, nil | ||
| } | ||
| return &http.Response{ | ||
| Request: r, | ||
| StatusCode: http.StatusOK, | ||
| Body: io.NopCloser(strings.NewReader(`{"permissions": ["can_run_queries"]}`)), | ||
| }, nil | ||
| }), | ||
| } | ||
| client, err := c.NewApiClient() | ||
| require.NoError(t, err) | ||
| ctx := context.Background() | ||
| var res map[string][]string | ||
| err = client.Do(ctx, "GET", "/api/2.0/permissions/object/id", httpclient.WithResponseUnmarshal(&res)) | ||
| assert.NoError(t, err) | ||
| assert.Equal(t, map[string][]string{"permissions": {"can_run_queries"}}, res) | ||
| } | ||
|
Comment on lines
+1
to
+54
There was a problem hiding this comment. I'd recommend structuring the unit tests differently by having this test focused on how This test could look like the following (I did not verify that the code works): type mock struct {
MaxFails int // number of times the failed Response is returned
FailResponse *http.Response // response to return in case of fail
FailError error // error to return in case of fail
NumCalls int // total number of calls
}
func (m *mock) RoundTrip(r *http.Request) (*http.Response, error) {
m.NumCalls++
if m.NumCalls <= m.MaxFails {
return m.FailResponse, n.FailError
}
return &http.Response{
Request: r,
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`{}`)),
}, nil
}
func (m *mock) SkipRetryOnIO() bool {
return true
}
func TestApiClient_Do_retries(t *testing.T) {
testCases := []struct{
desc string
config *Config
errorRetrier ErrorRetrier
wantNumCalls int
} {
{
desc: "nil retrier",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
wantNumCalls: 1,
},
{
desc: "no retry",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(context.Context, *http.Request, *common.ResponseWrapper, error) bool {
return false
},
wantNumCalls: 1,
},
{
desc: "retry 1 time",
mock: &mock{
MaxFails: 1,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(context.Context, *http.Request, *common.ResponseWrapper, error) bool {
return true
},
wantNumCalls: 2,
},
{
desc: "retry 2 times",
mock: &mock{
MaxFails: 2,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(_ context.Context, _ *http.Request, _ *common.ResponseWrapper, _ error) bool {
return true
},
wantNumCalls: 3,
},
{
desc: "retry 3 times",
mock: &mock{
MaxFails: 3,
FailResponse: &http.Response{StatusCode: http.StatusGatewayTimeout}
}
errorRetrier: func(_ context.Context, _ *http.Request, _ *common.ResponseWrapper, _ error) bool {
return true
},
wantNumCalls: 4,
},
}
func _, tc := range testCases {
t.Run(tc.desc, func(t *testing.T) {
cfg := &Config{HTTPTransport: tc.mock}
client, err := cfg.NewApiClient()
client.ErrorRetrier = tc.errorRetrier
err = client.Do(context.Background(), "GET", "test-path")
gotNumCalls = tc.mock.NumCalls
if gotNumCalls != tc.wantNumCalls {
t.Errorf("got %d calls, want %d", gotNumCalls, tc.wantNumCalls)
}
})
}
}Please feel free to ignore this comment if this is too much work or if the ApiClient cannot be instrumented that easily. There was a problem hiding this comment. It definitely can be instrumented this way, and this is a nice test case to use (I'll adapt it and include it in this PR). However, I did want to specifically test the get permissions pathway. Essentially, this tests that "the client returned by Config.GetApiClient() correctly implements retry on 504." I will add more test cases here though. There was a problem hiding this comment.
Sounds good to me as long as this complements the overall testing of the retry logic. |
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can add this back in if desired, just a small formatting change.