226 lines
8.0 KiB
Go
226 lines
8.0 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
)
|
|
|
|
func TestPreparedGenerationFreezesPayloadAndAuthentication(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
initialKey string
|
|
wantAuth string
|
|
}{
|
|
{name: "absent optional key"},
|
|
{name: "present key", initialKey: "original-key", wantAuth: "Bearer original-key"},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
const envName = "PROMPTKIT_PREPARED_GENERATION_TEST_KEY"
|
|
t.Setenv(envName, tc.initialKey)
|
|
var bodies, auth, urls []string
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://example.invalid/v1", Model: "default-model",
|
|
HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
|
body, err := io.ReadAll(req.Body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
bodies = append(bodies, string(body))
|
|
auth = append(auth, req.Header.Get("Authorization"))
|
|
urls = append(urls, req.URL.String())
|
|
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`))}, nil
|
|
})},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req := domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "original"}}},
|
|
Target: domain.ExecutionTarget{APIKeyEnv: envName},
|
|
}
|
|
prepared, err := client.PrepareGeneration(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.Prompt.Messages[0].Content = "changed"
|
|
t.Setenv(envName, "later-key")
|
|
for range 2 {
|
|
outcome := prepared.Execute(context.Background())
|
|
if outcome.Err != nil || outcome.Response == nil || outcome.Response.Content != "ok" {
|
|
t.Fatalf("Execute = %+v", outcome)
|
|
}
|
|
}
|
|
if len(bodies) != 2 || bodies[0] != bodies[1] || !strings.Contains(bodies[0], "original") || strings.Contains(bodies[0], "changed") || !reflect.DeepEqual(auth, []string{tc.wantAuth, tc.wantAuth}) || !reflect.DeepEqual(urls, []string{"https://example.invalid/v1/chat/completions", "https://example.invalid/v1/chat/completions"}) {
|
|
t.Fatalf("bodies=%v auth=%v urls=%v", bodies, auth, urls)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPreparedGenerationAttemptsHaveIndependentContextsAndCloseBodies(t *testing.T) {
|
|
var contexts []context.Context
|
|
var closed int
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://example.invalid/v1", Model: "model",
|
|
HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
|
if err := req.Context().Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
contexts = append(contexts, req.Context())
|
|
return &http.Response{StatusCode: http.StatusOK, Body: &countingBody{
|
|
Reader: strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`),
|
|
closed: &closed,
|
|
}}, nil
|
|
})},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
|
|
Target: domain.ExecutionTarget{TimeoutSeconds: 1},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for range 2 {
|
|
if outcome := prepared.Execute(context.Background()); outcome.Err != nil {
|
|
t.Fatal(outcome.Err)
|
|
}
|
|
}
|
|
if len(contexts) != 2 || contexts[0] == contexts[1] || closed != 2 {
|
|
t.Fatalf("contexts=%d, closed=%d", len(contexts), closed)
|
|
}
|
|
for _, ctx := range contexts {
|
|
if _, ok := ctx.Deadline(); !ok || !errors.Is(ctx.Err(), context.Canceled) {
|
|
t.Fatalf("attempt context deadline/error = %v/%v", ok, ctx.Err())
|
|
}
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if outcome := prepared.Execute(ctx); !errors.Is(outcome.Err, context.Canceled) {
|
|
t.Fatalf("canceled attempt = %+v", outcome)
|
|
}
|
|
if outcome := prepared.Execute(context.Background()); outcome.Err != nil || outcome.Response == nil {
|
|
t.Fatalf("later attempt = %+v", outcome)
|
|
}
|
|
}
|
|
|
|
func TestPreparedGenerationRetainsUsageFromMalformedChoice(t *testing.T) {
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://example.invalid/v1", Model: "model",
|
|
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":0},"choices":[]}`))}, nil
|
|
})},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
outcome := prepared.Execute(context.Background())
|
|
if !errors.Is(outcome.Err, ErrMalformedResponse) || outcome.Response != nil || outcome.UsagePresence == nil || !outcome.UsagePresence.PromptTokens || outcome.Usage.PromptTokens != 0 {
|
|
t.Fatalf("malformed choice = %+v", outcome)
|
|
}
|
|
}
|
|
|
|
func TestPreparedGenerationDoesNotUseInvalidlyFramedUsage(t *testing.T) {
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://example.invalid/v1", Model: "model",
|
|
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(`{"usage":{"prompt_tokens":9},"choices":[]}{}`))}, nil
|
|
})},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
outcome := prepared.Execute(context.Background())
|
|
if !errors.Is(outcome.Err, ErrMalformedResponse) || outcome.UsagePresence != nil || outcome.Usage != (domain.TokenUsage{}) {
|
|
t.Fatalf("invalid framing = %+v", outcome)
|
|
}
|
|
}
|
|
|
|
func TestPreparedGenerationRetainsReadCancellation(t *testing.T) {
|
|
closed := false
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{
|
|
BaseURL: "https://example.invalid/v1", Model: "model",
|
|
HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{StatusCode: http.StatusOK, Body: &readErrorBody{err: context.Canceled, closed: &closed}}, nil
|
|
})},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
prepared, err := client.PrepareGeneration(domain.GenerateRequest{
|
|
Prompt: domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: "hi"}}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
outcome := prepared.Execute(context.Background())
|
|
if !errors.Is(outcome.Err, context.Canceled) || !errors.Is(outcome.Err, ErrTransport) || errors.Is(outcome.Err, ErrMalformedResponse) || !closed {
|
|
t.Fatalf("read cancellation = %+v, closed=%t", outcome, closed)
|
|
}
|
|
}
|
|
|
|
func TestRetryNotBefore(t *testing.T) {
|
|
received := time.Date(2026, 9, 23, 12, 0, 0, 0, time.UTC)
|
|
for _, tc := range []struct {
|
|
name string
|
|
value string
|
|
want time.Time
|
|
}{
|
|
{name: "seconds", value: " 12 ", want: received.Add(12 * time.Second)},
|
|
{name: "date", value: "Wed, 23 Sep 2026 12:00:15 GMT", want: received.Add(15 * time.Second)},
|
|
{name: "past date", value: "Wed, 23 Sep 2026 11:59:59 GMT", want: received},
|
|
{name: "negative", value: "-2"},
|
|
{name: "fractional", value: "1.5"},
|
|
{name: "malformed", value: "tomorrow"},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if got := retryNotBefore([]string{tc.value, "999"}, received); !got.Equal(tc.want) {
|
|
t.Fatalf("not before = %v, want %v", got, tc.want)
|
|
}
|
|
})
|
|
}
|
|
if got := retryNotBefore([]string{"999999999999999999999999999999999999"}, received); !got.After(received) {
|
|
t.Fatalf("large valid delay wrapped: %v", got)
|
|
}
|
|
}
|
|
|
|
type countingBody struct {
|
|
io.Reader
|
|
closed *int
|
|
}
|
|
|
|
func (b *countingBody) Close() error {
|
|
*b.closed++
|
|
return nil
|
|
}
|
|
|
|
type readErrorBody struct {
|
|
err error
|
|
closed *bool
|
|
}
|
|
|
|
func (b *readErrorBody) Read([]byte) (int, error) { return 0, b.err }
|
|
func (b *readErrorBody) Close() error { *b.closed = true; return nil }
|