84 lines
2.4 KiB
Go
84 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type runnerFunc func(context.Context, string, json.RawMessage) (json.RawMessage, error)
|
|
|
|
func (f runnerFunc) run(ctx context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
|
|
return f(ctx, model, body)
|
|
}
|
|
|
|
func TestHTTPStream(t *testing.T) {
|
|
for _, failed := range []bool{false, true} {
|
|
server := httptest.NewServer(handler{heartbeat: time.Millisecond, runner: runnerFunc(func(_ context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
|
|
if model != "partner-model@001" || string(body) != `{"contents":[]}` {
|
|
t.Errorf("changed request %s %s", model, body)
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
if failed {
|
|
return nil, errors.New("synthetic failure")
|
|
}
|
|
return json.RawMessage(`{"candidates":[{"finishReason":"STOP"}]}`), nil
|
|
})})
|
|
req, _ := http.NewRequest("POST", server.URL+"/v1/chat/completions", strings.NewReader(`{"contents":[]}`))
|
|
req.Header.Set("x-batch-model", "partner-model@001")
|
|
req.Header.Set("x-batch-stream", "true")
|
|
response, err := server.Client().Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
data, err := io.ReadAll(response.Body)
|
|
_ = response.Body.Close()
|
|
server.Close()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.Contains(string(data), ": batch pending") || strings.Contains(string(data), "data:") == failed {
|
|
t.Fatalf("failed=%v body=%s", failed, data)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCleanupWait(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
cancelled := make(chan struct{})
|
|
cleanup := make(chan struct{})
|
|
finished := make(chan struct{})
|
|
h := handler{runner: runnerFunc(func(ctx context.Context, _ string, _ json.RawMessage) (json.RawMessage, error) {
|
|
<-ctx.Done()
|
|
close(cancelled)
|
|
<-cleanup
|
|
return nil, ctx.Err()
|
|
})}
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"contents":[]}`)).WithContext(ctx)
|
|
req.Header.Set("x-batch-model", "partner-model@001")
|
|
req.Header.Set("x-batch-stream", "false")
|
|
go func() {
|
|
h.ServeHTTP(httptest.NewRecorder(), req)
|
|
close(finished)
|
|
}()
|
|
cancel()
|
|
<-cancelled
|
|
select {
|
|
case <-finished:
|
|
t.Error("handler returned before cloud cleanup")
|
|
case <-time.After(20 * time.Millisecond):
|
|
}
|
|
close(cleanup)
|
|
select {
|
|
case <-finished:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("handler did not finish after cleanup")
|
|
}
|
|
}
|