143 lines · 4.2 KB · Go
RawDownload
package executor import (	"context"	"net/http"	"net/http/httptest"	"testing"	"time" 	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator") func TestOpenAICompatRetryAfter(t *testing.T) {	now := time.Date(2026, time.September, 3, 12, 0, 0, 0, time.UTC)	tests := []struct {		name    string		status  int		headers http.Header		body    string		want    *time.Duration	}{		{			name:    "delta seconds header",			status:  http.StatusTooManyRequests,			headers: http.Header{"Retry-After": {"17"}},			want:    durationPointer(17 * time.Second),		},		{			name:    "http date header",			status:  http.StatusTooManyRequests,			headers: http.Header{"Retry-After": {now.Add(23 * time.Second).Format(http.TimeFormat)}},			want:    durationPointer(23 * time.Second),		},		{			name:   "explicit TPM code fallback",			status: http.StatusTooManyRequests,			body:   `{"error":{"code":"ModelAccountTpmRateLimitExceeded","message":"TPM limit exceeded"}}`,			want:   durationPointer(time.Minute),		},		{			name:   "TPM message fallback",			status: http.StatusTooManyRequests,			body:   `{"error":{"message":"TPM (Tokens Per Minute) limit of this model is exceeded"}}`,			want:   durationPointer(time.Minute),		},		{			name:    "provider header wins over fallback",			status:  http.StatusTooManyRequests,			headers: http.Header{"Retry-After": {"5"}},			body:    `{"error":{"code":"ModelAccountTpmRateLimitExceeded"}}`,			want:    durationPointer(5 * time.Second),		},		{			name:   "generic 429 has no invented deadline",			status: http.StatusTooManyRequests,			body:   `{"error":{"code":"rate_limit"}}`,		},		{			name:    "non-429 ignores header",			status:  http.StatusServiceUnavailable,			headers: http.Header{"Retry-After": {"30"}},		},	} 	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			got := openAICompatRetryAfter(test.status, test.headers, []byte(test.body), now)			if test.want == nil {				if got != nil {					t.Fatalf("retry-after = %v, want nil", *got)				}				return			}			if got == nil || *got != *test.want {				t.Fatalf("retry-after = %v, want %v", got, *test.want)			}		})	}} func TestOpenAICompatExecutorPropagatesRetryAfter(t *testing.T) {	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {		w.Header().Set("Content-Type", "application/json")		w.Header().Set("Retry-After", "7")		w.WriteHeader(http.StatusTooManyRequests)		_, _ = w.Write([]byte(`{"error":{"code":"rate_limit","message":"try later"}}`))	}))	t.Cleanup(server.Close) 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	request := cliproxyexecutor.Request{		Model:   "compatible-model",		Payload: []byte(`{"model":"compatible-model","messages":[{"role":"user","content":"hi"}]}`),	}	tests := []struct {		name   string		invoke func() error	}{		{			name: "nonstream",			invoke: func() error {				_, errExecute := executor.Execute(context.Background(), auth, request, cliproxyexecutor.Options{					SourceFormat: sdktranslator.FromString("openai"),				})				return errExecute			},		},		{			name: "stream bootstrap",			invoke: func() error {				_, errExecute := executor.ExecuteStream(context.Background(), auth, request, cliproxyexecutor.Options{					SourceFormat: sdktranslator.FromString("openai"),					Stream:       true,				})				return errExecute			},		},	} 	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			errExecute := test.invoke()			if errExecute == nil {				t.Fatal("expected rate-limit error")			}			retryable, ok := errExecute.(interface{ RetryAfter() *time.Duration })			if !ok || retryable.RetryAfter() == nil || *retryable.RetryAfter() != 7*time.Second {				t.Fatalf("retry-after = %v, want 7s", retryable)			}		})	}} func durationPointer(value time.Duration) *time.Duration {	return &value}