279 lines
9.2 KiB
Go
279 lines
9.2 KiB
Go
package generation
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
|
|
)
|
|
|
|
type transportFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f transportFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }
|
|
|
|
func response(status int, body string) *http.Response {
|
|
return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}
|
|
}
|
|
|
|
func testRequest(retries uint64) domain.GenerateRequest {
|
|
return domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hello"}}},
|
|
Target: domain.ExecutionTarget{BackendID: "limited", Endpoint: "https://example.invalid/v1", Model: "model"},
|
|
TransientRetries: retries,
|
|
}
|
|
}
|
|
|
|
func testCoordinator(t *testing.T, transport transportFunc) *Coordinator {
|
|
t.Helper()
|
|
manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
|
|
"limited": {ConcurrencyLimit: 1, QueueCapacity: 1},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client, err := llm.NewOpenAICompatibleClient(llm.OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: transport}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
coordinator := NewBuiltIn(manager, client)
|
|
coordinator.jitter = func(time.Duration) time.Duration { return 0 }
|
|
coordinator.wait = func(ctx context.Context, _ time.Duration) error { return ctx.Err() }
|
|
return coordinator
|
|
}
|
|
|
|
func TestBuiltInRetryEligibilityAndExhaustion(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
status int
|
|
eligible bool
|
|
}{
|
|
{status: 429, eligible: true}, {status: 500, eligible: true},
|
|
{status: 502, eligible: true}, {status: 503, eligible: true},
|
|
{status: 504, eligible: true}, {status: 400}, {status: 501},
|
|
} {
|
|
t.Run(http.StatusText(tc.status), func(t *testing.T) {
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
return response(tc.status, `{"error":{}}`), nil
|
|
}
|
|
return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
|
|
})
|
|
result, err := coordinator.Generate(context.Background(), testRequest(1))
|
|
if tc.eligible {
|
|
if err != nil || result == nil || result.Content != "ok" || calls != 2 {
|
|
t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
|
|
}
|
|
} else if result != nil || err == nil || calls != 1 {
|
|
t.Fatalf("result=%+v, error=%v, calls=%d", result, err, calls)
|
|
}
|
|
})
|
|
}
|
|
|
|
for _, budget := range []uint64{0, 2} {
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return response(503, `{"error":{}}`), nil
|
|
})
|
|
_, err := coordinator.Generate(context.Background(), testRequest(budget))
|
|
var exhausted *ExhaustedError
|
|
if !errors.Is(err, llm.ErrUnexpectedStatus) || calls != int(budget+1) || errors.As(err, &exhausted) != (budget > 0) {
|
|
t.Fatalf("budget=%d calls=%d error=%v", budget, calls, err)
|
|
}
|
|
if budget > 0 && exhausted.Attempts() != budget+1 {
|
|
t.Fatalf("attempts=%d", exhausted.Attempts())
|
|
}
|
|
}
|
|
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
return response(503, `{"error":{}}`), nil
|
|
}
|
|
return response(400, `{"error":{}}`), nil
|
|
})
|
|
_, err := coordinator.Generate(context.Background(), testRequest(3))
|
|
var exhausted *ExhaustedError
|
|
if calls != 2 || !errors.Is(err, llm.ErrUnexpectedStatus) || errors.As(err, &exhausted) {
|
|
t.Fatalf("later permanent failure: calls=%d, error=%v", calls, err)
|
|
}
|
|
}
|
|
|
|
func TestBuiltInRetryWaitHonorsContextAndProviderMinimum(t *testing.T) {
|
|
when := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
|
|
if got := coordinator.backoff(1); got != 250*time.Millisecond {
|
|
t.Fatalf("first backoff=%v", got)
|
|
}
|
|
if got := coordinator.backoff(math.MaxUint64); got != 4*time.Second {
|
|
t.Fatalf("capped backoff=%v", got)
|
|
}
|
|
coordinator.jitter = func(max time.Duration) time.Duration { return max }
|
|
if got := coordinator.backoff(math.MaxUint64); got != 8*time.Second {
|
|
t.Fatalf("maximum equal jitter=%v", got)
|
|
}
|
|
if got := retryDelay(500*time.Millisecond, when.Add(3*time.Second), when); got != 3*time.Second {
|
|
t.Fatalf("provider minimum=%v", got)
|
|
}
|
|
if got := retryDelay(500*time.Millisecond, when.Add(-time.Second), when); got != 500*time.Millisecond {
|
|
t.Fatalf("past provider minimum=%v", got)
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
calls := 0
|
|
coordinator = testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return response(503, `{"error":{}}`), nil
|
|
})
|
|
coordinator.wait = func(ctx context.Context, delay time.Duration) error {
|
|
if delay < 250*time.Millisecond {
|
|
t.Fatalf("delay=%v", delay)
|
|
}
|
|
cancel()
|
|
return ctx.Err()
|
|
}
|
|
_, err := coordinator.Generate(ctx, testRequest(2))
|
|
if !errors.Is(err, context.Canceled) || calls != 1 {
|
|
t.Fatalf("canceled wait: error=%v calls=%d", err, calls)
|
|
}
|
|
}
|
|
|
|
func TestBuiltInRetryUsesSeparateAttemptDeadlines(t *testing.T) {
|
|
var contexts []context.Context
|
|
coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
|
|
contexts = append(contexts, req.Context())
|
|
if len(contexts) == 1 {
|
|
return response(503, `{"error":{}}`), nil
|
|
}
|
|
return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
|
|
})
|
|
req := testRequest(1)
|
|
req.Target.TimeoutSeconds = 1
|
|
if _, err := coordinator.Generate(context.Background(), req); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(contexts) != 2 || contexts[0] == contexts[1] {
|
|
t.Fatalf("attempt contexts=%d", len(contexts))
|
|
}
|
|
for _, ctx := range contexts {
|
|
if _, ok := ctx.Deadline(); !ok || !errors.Is(ctx.Err(), context.Canceled) {
|
|
t.Fatalf("attempt deadline/error=%t/%v", ok, ctx.Err())
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCanceledCapacityWaitDispatchesNothing(t *testing.T) {
|
|
calls := 0
|
|
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
|
|
calls++
|
|
return response(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
|
|
})
|
|
release, err := coordinator.manager.Acquire(context.Background(), "limited")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer release()
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
result := make(chan error, 1)
|
|
go func() { _, err := coordinator.Generate(ctx, testRequest(2)); result <- err }()
|
|
cancel()
|
|
if err := <-result; !errors.Is(err, context.Canceled) || calls != 0 {
|
|
t.Fatalf("capacity cancellation: error=%v calls=%d", err, calls)
|
|
}
|
|
}
|
|
|
|
func TestBuiltInRetryFreesPermitAndFreezesCredential(t *testing.T) {
|
|
const envName = "PROMPTKIT_GENERATION_RETRY_TEST_KEY"
|
|
t.Setenv(envName, "first-key")
|
|
var calls atomic.Int64
|
|
var firstAuth, retryAuth string
|
|
coordinator := testCoordinator(t, func(req *http.Request) (*http.Response, error) {
|
|
switch calls.Add(1) {
|
|
case 1:
|
|
firstAuth = req.Header.Get("Authorization")
|
|
return response(503, `{"error":{}}`), nil
|
|
case 2:
|
|
return response(200, `{"choices":[{"message":{"content":"other"}}]}`), nil
|
|
default:
|
|
retryAuth = req.Header.Get("Authorization")
|
|
return response(200, `{"choices":[{"message":{"content":"retry"}}]}`), nil
|
|
}
|
|
})
|
|
entered := make(chan struct{})
|
|
resume := make(chan struct{})
|
|
coordinator.wait = func(ctx context.Context, _ time.Duration) error {
|
|
close(entered)
|
|
select {
|
|
case <-resume:
|
|
return ctx.Err()
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
req := testRequest(1)
|
|
req.Target.APIKeyEnv = envName
|
|
first := make(chan error, 1)
|
|
go func() { _, err := coordinator.Generate(context.Background(), req); first <- err }()
|
|
<-entered
|
|
t.Setenv(envName, "later-key")
|
|
other, err := coordinator.Generate(context.Background(), testRequest(0))
|
|
if err != nil || other == nil || other.Content != "other" {
|
|
t.Fatalf("other call while backoff blocked = %+v, %v", other, err)
|
|
}
|
|
close(resume)
|
|
if err := <-first; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if calls.Load() != 3 || firstAuth != "Bearer first-key" || retryAuth != firstAuth {
|
|
t.Fatalf("calls=%d, first auth=%q, retry auth=%q", calls.Load(), firstAuth, retryAuth)
|
|
}
|
|
}
|
|
|
|
func TestInjectedClientDispatchesOnceAndReleasesAfterPanic(t *testing.T) {
|
|
manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{"limited": {ConcurrencyLimit: 1}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client := &injectedClient{panicFirst: true}
|
|
coordinator := NewInjected(manager, client)
|
|
func() {
|
|
defer func() {
|
|
if recover() == nil {
|
|
t.Fatal("expected injected panic")
|
|
}
|
|
}()
|
|
_, _ = coordinator.Generate(context.Background(), testRequest(5))
|
|
}()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
result, err := coordinator.Generate(ctx, testRequest(5))
|
|
if err != nil || result == nil || client.calls != 2 {
|
|
t.Fatalf("after panic result=%+v error=%v calls=%d", result, err, client.calls)
|
|
}
|
|
}
|
|
|
|
type injectedClient struct {
|
|
calls int
|
|
panicFirst bool
|
|
}
|
|
|
|
func (c *injectedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
|
c.calls++
|
|
if c.panicFirst && c.calls == 1 {
|
|
panic("injected panic")
|
|
}
|
|
return &domain.GenerateResponse{Content: "ok"}, nil
|
|
}
|