212 lines
6.0 KiB
Go
212 lines
6.0 KiB
Go
package zarinhub
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRetryTransportPreservesAndRewindsBody(t *testing.T) {
|
|
var attempts int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
att := atomic.AddInt32(&attempts, 1)
|
|
body, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
t.Errorf("read body error: %v", err)
|
|
}
|
|
if string(body) != `{"payload":"important"}` {
|
|
t.Errorf("unexpected body on attempt %d: %q", att, string(body))
|
|
}
|
|
if att < 3 {
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{"status":"ok"}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
transport := NewRetryTransport(server.Client().Transport, time.Second, 3)
|
|
transport.retryWait = func(ctx context.Context, delay time.Duration) bool {
|
|
return true
|
|
}
|
|
client := &http.Client{Transport: transport}
|
|
|
|
req, err := http.NewRequestWithContext(
|
|
context.Background(),
|
|
http.MethodPost,
|
|
server.URL,
|
|
bytes.NewReader([]byte(`{"payload":"important"}`)),
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("NewRequest error: %v", err)
|
|
}
|
|
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("client.Do error: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK)
|
|
}
|
|
if attempts != 3 {
|
|
t.Errorf("attempts = %d, want 3", attempts)
|
|
}
|
|
}
|
|
|
|
func TestRetryTransportDoesNotRetryClientErrors(t *testing.T) {
|
|
tests := []int{
|
|
http.StatusBadRequest,
|
|
http.StatusUnauthorized,
|
|
http.StatusForbidden,
|
|
http.StatusNotFound,
|
|
http.StatusUnprocessableEntity,
|
|
}
|
|
|
|
for _, statusCode := range tests {
|
|
t.Run(fmt.Sprintf("status_%d", statusCode), func(t *testing.T) {
|
|
var attempts int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
atomic.AddInt32(&attempts, 1)
|
|
w.WriteHeader(statusCode)
|
|
}))
|
|
defer server.Close()
|
|
|
|
transport := NewRetryTransport(server.Client().Transport, time.Second, 2)
|
|
client := &http.Client{Transport: transport}
|
|
|
|
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, nil)
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("client.Do error: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != statusCode {
|
|
t.Errorf("statusCode = %d, want %d", resp.StatusCode, statusCode)
|
|
}
|
|
if attempts != 1 {
|
|
t.Errorf("attempts = %d, want 1", attempts)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRetryTransportPerAttemptTimeout(t *testing.T) {
|
|
var attempts int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
att := atomic.AddInt32(&attempts, 1)
|
|
if att == 1 {
|
|
time.Sleep(150 * time.Millisecond)
|
|
w.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer server.Close()
|
|
|
|
transport := NewRetryTransport(server.Client().Transport, 40*time.Millisecond, 2)
|
|
transport.retryWait = func(ctx context.Context, delay time.Duration) bool {
|
|
return true
|
|
}
|
|
client := &http.Client{Transport: transport}
|
|
|
|
// Parent context has plenty of time (5s)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("client.Do error: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Errorf("status = %d, want 200", resp.StatusCode)
|
|
}
|
|
if attempts != 2 {
|
|
t.Errorf("attempts = %d, want 2", attempts)
|
|
}
|
|
}
|
|
|
|
func TestRetryTransportParentContextCancellation(t *testing.T) {
|
|
var attempts int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
atomic.AddInt32(&attempts, 1)
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}))
|
|
defer server.Close()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
transport := NewRetryTransport(server.Client().Transport, time.Second, 3)
|
|
transport.retryWait = func(c context.Context, delay time.Duration) bool {
|
|
cancel() // cancel parent context during backoff
|
|
return waitForRetry(ctx, delay)
|
|
}
|
|
client := &http.Client{Transport: transport}
|
|
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
|
|
_, err := client.Do(req)
|
|
if err == nil {
|
|
t.Fatal("expected error, got nil")
|
|
}
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Errorf("error = %v, want context.Canceled", err)
|
|
}
|
|
if attempts != 1 {
|
|
t.Errorf("attempts = %d, want 1", attempts)
|
|
}
|
|
}
|
|
|
|
func TestRetryTransportLoggingFormat(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}))
|
|
defer server.Close()
|
|
|
|
var buf bytes.Buffer
|
|
logger := slog.New(slog.NewJSONHandler(&buf, nil))
|
|
|
|
transport := NewRetryTransport(server.Client().Transport, time.Second, 1, logger)
|
|
transport.retryWait = func(ctx context.Context, delay time.Duration) bool {
|
|
return true
|
|
}
|
|
client := &http.Client{Transport: transport}
|
|
|
|
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, nil)
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
logStr := buf.String()
|
|
if !strings.Contains(logStr, "[INITIAL] zarinhub request") {
|
|
t.Errorf("expected logs to contain '[INITIAL] zarinhub request', got: %s", logStr)
|
|
}
|
|
if !strings.Contains(logStr, `"is_retry":false`) {
|
|
t.Errorf("expected logs to contain '\"is_retry\":false', got: %s", logStr)
|
|
}
|
|
if !strings.Contains(logStr, "[RETRY] zarinhub request attempt 1") {
|
|
t.Errorf("expected logs to contain '[RETRY] zarinhub request attempt 1', got: %s", logStr)
|
|
}
|
|
if !strings.Contains(logStr, "[RETRY] zarinhub request failed, retrying in") {
|
|
t.Errorf("expected logs to contain retry warning, got: %s", logStr)
|
|
}
|
|
if !strings.Contains(logStr, "[RETRY] zarinhub request failed: all retries exhausted") {
|
|
t.Errorf("expected logs to contain exhausted message, got: %s", logStr)
|
|
}
|
|
}
|