Skip to content

Commit 3916f75

Browse files
somanshreddysomanshreddy
committed
client: add retry with exponential backoff (PRINFRA-122)
1 parent 1de7729 commit 3916f75

6 files changed

Lines changed: 838 additions & 2 deletions

File tree

‎cmd/heygen/root.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@ func newRootCmd(version string, formatter output.Formatter) *cobra.Command {
1717
root := &cobra.Command{
1818
Use: "heygen",
1919
Short: "HeyGen CLI — create and manage videos, avatars, and more",
20-
// NOTE: env var list is hardcoded. Keep in sync with envVarByKey in env_provider.go.
20+
// NOTE: env var list is hardcoded. Keep config-related entries in sync with
21+
// envVarByKey in env_provider.go and update operational settings here too.
2122
Long: `HeyGen CLI — create and manage videos, avatars, and more.
2223
2324
Environment Variables:
2425
HEYGEN_API_KEY API key for authentication (overrides stored credentials)
2526
HEYGEN_OUTPUT Output format: json, human (default: json)
2627
HEYGEN_NO_ANALYTICS Disable analytics when set (default: enabled)
28+
HEYGEN_MAX_RETRIES Max retries for transient errors (default: 2, 0 to disable)
2729
HEYGEN_CONFIG_DIR Override config directory (default: ~/.heygen)`,
2830
Version: version,
2931
SilenceUsage: true, // we handle usage errors ourselves

‎internal/client/client.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ type Client struct {
1919
baseURL string
2020
apiKey string
2121
userAgent string
22+
retry RetryConfig
2223
}
2324

2425
// Option configures the Client.
@@ -34,6 +35,16 @@ func WithHTTPClient(hc *http.Client) Option {
3435
return func(c *Client) { c.httpClient = hc }
3536
}
3637

38+
// WithRetry overrides the default retry configuration.
39+
func WithRetry(config RetryConfig) Option {
40+
return func(c *Client) { c.retry = config }
41+
}
42+
43+
// WithNoRetry disables retries entirely.
44+
func WithNoRetry() Option {
45+
return func(c *Client) { c.retry.MaxRetries = 0 }
46+
}
47+
3748
// WithUserAgent overrides the default User-Agent header.
3849
func WithUserAgent(ua string) Option {
3950
return func(c *Client) { c.userAgent = ua }
@@ -46,10 +57,23 @@ func New(apiKey string, opts ...Option) *Client {
4657
baseURL: DefaultBaseURL,
4758
apiKey: apiKey,
4859
userAgent: DefaultUserAgent,
60+
retry: DefaultRetryConfig(),
4961
}
5062
for _, opt := range opts {
5163
opt(c)
5264
}
65+
66+
copied := *c.httpClient
67+
transport := copied.Transport
68+
if transport == nil {
69+
transport = http.DefaultTransport
70+
}
71+
copied.Transport = &retryTransport{
72+
base: transport,
73+
config: c.retry,
74+
}
75+
c.httpClient = &copied
76+
5377
return c
5478
}
5579

‎internal/client/client_test.go‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"net/http"
55
"net/http/httptest"
66
"testing"
7+
"time"
78
)
89

910
func TestClient_Do_InjectsHeaders(t *testing.T) {
@@ -65,4 +66,92 @@ func TestClient_Defaults(t *testing.T) {
6566
if c.userAgent != DefaultUserAgent {
6667
t.Errorf("userAgent = %q, want %q", c.userAgent, DefaultUserAgent)
6768
}
69+
if c.retry.MaxRetries != 2 {
70+
t.Errorf("retry.MaxRetries = %d, want %d", c.retry.MaxRetries, 2)
71+
}
72+
if c.retry.BaseDelay != time.Second {
73+
t.Errorf("retry.BaseDelay = %v, want %v", c.retry.BaseDelay, time.Second)
74+
}
75+
if c.retry.MaxDelay != 30*time.Second {
76+
t.Errorf("retry.MaxDelay = %v, want %v", c.retry.MaxDelay, 30*time.Second)
77+
}
78+
if _, ok := c.httpClient.Transport.(*retryTransport); !ok {
79+
t.Fatalf("httpClient.Transport = %T, want *retryTransport", c.httpClient.Transport)
80+
}
81+
}
82+
83+
func TestClient_WithNoRetry(t *testing.T) {
84+
c := New("key", WithNoRetry())
85+
if c.retry.MaxRetries != 0 {
86+
t.Errorf("retry.MaxRetries = %d, want %d", c.retry.MaxRetries, 0)
87+
}
88+
}
89+
90+
func TestClient_WithHTTPClientPreservesTimeout(t *testing.T) {
91+
custom := &http.Client{
92+
Timeout: 5 * time.Second,
93+
Transport: &stubTransport{},
94+
}
95+
96+
c := New("key", WithHTTPClient(custom))
97+
if c.httpClient.Timeout != 5*time.Second {
98+
t.Errorf("httpClient.Timeout = %v, want %v", c.httpClient.Timeout, 5*time.Second)
99+
}
100+
rt, ok := c.httpClient.Transport.(*retryTransport)
101+
if !ok {
102+
t.Fatalf("httpClient.Transport = %T, want *retryTransport", c.httpClient.Transport)
103+
}
104+
if _, ok := rt.base.(*stubTransport); !ok {
105+
t.Fatalf("retryTransport.base = %T, want *stubTransport", rt.base)
106+
}
107+
}
108+
109+
func TestClient_WithHTTPClientNoMutation(t *testing.T) {
110+
base := &stubTransport{}
111+
custom := &http.Client{
112+
Timeout: 5 * time.Second,
113+
Transport: base,
114+
}
115+
116+
c1 := New("key", WithHTTPClient(custom))
117+
c2 := New("key", WithHTTPClient(custom))
118+
119+
if custom.Transport != base {
120+
t.Fatalf("custom.Transport mutated to %T, want original transport", custom.Transport)
121+
}
122+
123+
rt1, ok := c1.httpClient.Transport.(*retryTransport)
124+
if !ok {
125+
t.Fatalf("c1.httpClient.Transport = %T, want *retryTransport", c1.httpClient.Transport)
126+
}
127+
rt2, ok := c2.httpClient.Transport.(*retryTransport)
128+
if !ok {
129+
t.Fatalf("c2.httpClient.Transport = %T, want *retryTransport", c2.httpClient.Transport)
130+
}
131+
if rt1.base != base {
132+
t.Fatalf("c1 retry base = %T, want original transport", rt1.base)
133+
}
134+
if rt2.base != base {
135+
t.Fatalf("c2 retry base = %T, want original transport", rt2.base)
136+
}
137+
}
138+
139+
func TestDefaultRetryConfig_FromEnv(t *testing.T) {
140+
t.Setenv("HEYGEN_MAX_RETRIES", "5")
141+
142+
cfg := DefaultRetryConfig()
143+
if cfg.MaxRetries != 5 {
144+
t.Errorf("cfg.MaxRetries = %d, want %d", cfg.MaxRetries, 5)
145+
}
146+
}
147+
148+
type stubTransport struct{}
149+
150+
func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
151+
return &http.Response{
152+
StatusCode: http.StatusOK,
153+
Body: http.NoBody,
154+
Header: make(http.Header),
155+
Request: req,
156+
}, nil
68157
}

‎internal/client/executor_test.go‎

Lines changed: 75 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ func TestExecute_NetworkError(t *testing.T) {
229229
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
230230
srv.Close()
231231

232-
c := New("key", WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
232+
c := New("key", WithBaseURL(srv.URL), WithHTTPClient(srv.Client()), WithNoRetry())
233233

234234
spec := &command.Spec{Endpoint: "/v3/videos", Method: "GET"}
235235
inv := &command.Invocation{PathParams: make(map[string]string), QueryParams: make(url.Values)}
@@ -245,6 +245,80 @@ func TestExecute_NetworkError(t *testing.T) {
245245
}
246246
}
247247

248+
func TestExecute_RetryOn429(t *testing.T) {
249+
var calls int
250+
251+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
252+
calls++
253+
if calls == 1 {
254+
w.WriteHeader(http.StatusTooManyRequests)
255+
_, _ = w.Write([]byte(`{"error":{"code":"rate_limited","message":"too many requests"}}`))
256+
return
257+
}
258+
w.WriteHeader(http.StatusOK)
259+
_, _ = w.Write([]byte(`{"data":[]}`))
260+
}))
261+
defer srv.Close()
262+
263+
c := New("key", WithBaseURL(srv.URL), WithHTTPClient(srv.Client()), WithRetry(RetryConfig{MaxRetries: 1}))
264+
265+
spec := &command.Spec{Endpoint: "/v3/videos", Method: "GET"}
266+
inv := &command.Invocation{PathParams: make(map[string]string), QueryParams: make(url.Values)}
267+
268+
result, err := c.Execute(spec, inv)
269+
if err != nil {
270+
t.Fatalf("unexpected error: %v", err)
271+
}
272+
if calls != 2 {
273+
t.Fatalf("calls = %d, want %d", calls, 2)
274+
}
275+
if string(result) != `{"data":[]}` {
276+
t.Fatalf("result = %s, want %s", result, `{"data":[]}`)
277+
}
278+
}
279+
280+
func TestExecute_RetryPreservesBody(t *testing.T) {
281+
var calls int
282+
var bodies []string
283+
284+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
285+
body, _ := io.ReadAll(r.Body)
286+
bodies = append(bodies, string(body))
287+
calls++
288+
if calls == 1 {
289+
w.WriteHeader(http.StatusTooManyRequests)
290+
_, _ = w.Write([]byte(`{"error":{"code":"rate_limited","message":"too many requests"}}`))
291+
return
292+
}
293+
w.WriteHeader(http.StatusOK)
294+
_, _ = w.Write([]byte(`{"id":"new123"}`))
295+
}))
296+
defer srv.Close()
297+
298+
c := New("key", WithBaseURL(srv.URL), WithHTTPClient(srv.Client()), WithRetry(RetryConfig{MaxRetries: 1}))
299+
300+
spec := &command.Spec{Endpoint: "/v3/videos", Method: "POST", BodyEncoding: "json"}
301+
inv := &command.Invocation{
302+
PathParams: make(map[string]string),
303+
QueryParams: make(url.Values),
304+
Body: map[string]any{"title": "My Video", "draft": true},
305+
}
306+
307+
result, err := c.Execute(spec, inv)
308+
if err != nil {
309+
t.Fatalf("unexpected error: %v", err)
310+
}
311+
if calls != 2 {
312+
t.Fatalf("calls = %d, want %d", calls, 2)
313+
}
314+
if len(bodies) != 2 || bodies[0] != bodies[1] {
315+
t.Fatalf("bodies = %#v, want two identical request bodies", bodies)
316+
}
317+
if string(result) != `{"id":"new123"}` {
318+
t.Fatalf("result = %s, want %s", result, `{"id":"new123"}`)
319+
}
320+
}
321+
248322
func TestExecute_MultipartUpload(t *testing.T) {
249323
var gotContentType string
250324
var gotFileContent string

0 commit comments

Comments
 (0)