merged from gramsrv upstream

internal/loadharness/rpc_retry_test.go

@@ -0,0 +1,63 @@
+package loadharness
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/iamxvbaba/td/tgerr"
+)
+
+func TestRPCWithFloodWaitPolicyRetriesTheSameCall(t *testing.T) {
+	attempts := 0
+	waits := make([]time.Duration, 0, 1)
+	policy := floodWaitPolicy{
+		maxRetries: 2,
+		maxWait:    5 * time.Second,
+		padding:    100 * time.Millisecond,
+		wait: func(_ context.Context, delay time.Duration) error {
+			waits = append(waits, delay)
+			return nil
+		},
+	}
+	result, err := rpcWithFloodWaitPolicy(context.Background(), time.Second, policy, func(context.Context) (int, error) {
+		attempts++
+		if attempts == 1 {
+			return 0, tgerr.New(420, "FLOOD_WAIT_2")
+		}
+		return 42, nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result != 42 || attempts != 2 {
+		t.Fatalf("result=%d attempts=%d", result, attempts)
+	}
+	if len(waits) != 1 || waits[0] != 2100*time.Millisecond {
+		t.Fatalf("waits=%v", waits)
+	}
+}
+
+func TestRPCWithFloodWaitPolicyDoesNotRetryOrdinaryError(t *testing.T) {
+	want := errors.New("boom")
+	attempts := 0
+	policy := floodWaitPolicy{maxRetries: 2, maxWait: time.Second, wait: waitForContext}
+	_, err := rpcWithFloodWaitPolicy(context.Background(), time.Second, policy, func(context.Context) (int, error) {
+		attempts++
+		return 0, want
+	})
+	if !errors.Is(err, want) || attempts != 1 {
+		t.Fatalf("err=%v attempts=%d", err, attempts)
+	}
+}
+
+func TestRPCWithFloodWaitPolicyRejectsExcessiveWait(t *testing.T) {
+	policy := floodWaitPolicy{maxRetries: 2, maxWait: time.Second, wait: waitForContext}
+	_, err := rpcWithFloodWaitPolicy(context.Background(), time.Second, policy, func(context.Context) (int, error) {
+		return 0, tgerr.New(420, "FLOOD_WAIT_2")
+	})
+	if err == nil {
+		t.Fatal("excessive FLOOD_WAIT unexpectedly retried")
+	}
+}