101 lines
3.9 KiB
Go
101 lines
3.9 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"testing/synctest"
|
|
"time"
|
|
|
|
"go.maximumdirect.net/promptkit"
|
|
)
|
|
|
|
func TestRunOperationTimeoutAttribution(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
}))
|
|
var terminal promptkit.ExecutionEvent
|
|
_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second,
|
|
Observer: func(event promptkit.ExecutionEvent) {
|
|
if event.Kind == promptkit.ExecutionFinished {
|
|
terminal = event
|
|
}
|
|
}})
|
|
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) || terminal.Failure != promptkit.ExecutionFailureDeadline {
|
|
t.Fatalf("error=%v terminal=%+v", err, terminal)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRunProviderDeadlineIsNotOperationTimeout(t *testing.T) {
|
|
engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
return nil, context.DeadlineExceeded
|
|
}))
|
|
_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Minute})
|
|
if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
|
|
t.Fatalf("provider deadline attribution: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestPreparedExecutionRetainsTimeoutWithoutIdleTimer(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
}))
|
|
req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}
|
|
handle, err := engine.PrepareExecution(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req.OperationTimeout = 0
|
|
time.Sleep(2 * time.Second)
|
|
start := time.Now()
|
|
_, err = engine.RunPrepared(context.Background(), handle)
|
|
if time.Since(start) != time.Second || !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) {
|
|
t.Fatalf("elapsed=%v error=%v", time.Since(start), err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestPreparedExecutionRespectsCallerDeadline(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
<-ctx.Done()
|
|
return nil, ctx.Err()
|
|
}))
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
handle, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: 5 * time.Second})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
time.Sleep(time.Second)
|
|
_, err = engine.RunPrepared(ctx, handle)
|
|
if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
|
|
t.Fatalf("caller deadline attribution: %v", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestPreparedExecutionDoesNotTimePreparation(t *testing.T) {
|
|
synctest.Test(t, func(t *testing.T) {
|
|
release := make(chan struct{})
|
|
reader := &capacityArtifactReader{entered: make(chan struct{}, 1), release: release}
|
|
engine := newBackendCapacityEngine(t, newCapacityGateClient(closedCapacityChannel(), 1), 1, capacityInt(0), reader)
|
|
req := capacityInputRequest("http://example.invalid/v1")
|
|
req.OperationTimeout = time.Second
|
|
go func() { time.Sleep(2 * time.Second); close(release) }()
|
|
start := time.Now()
|
|
handle, err := engine.PrepareExecution(context.Background(), req)
|
|
if err != nil || time.Since(start) != 2*time.Second {
|
|
t.Fatalf("prepare elapsed=%v error=%v", time.Since(start), err)
|
|
}
|
|
if _, err = engine.RunPrepared(context.Background(), handle); err != nil {
|
|
t.Fatalf("execution after preparation: %v", err)
|
|
}
|
|
})
|
|
}
|