82 lines
2.4 KiB
Go
82 lines
2.4 KiB
Go
package llm
|
|
|
|
// Observes how much of a gateway response the client will read. The body is read before
|
|
// the status is looked at, so a misbehaving gateway or proxy that streams forever was
|
|
// bounded only by the client's overall timeout (#53).
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// endlessGateway answers with the given status and streams filler until the client hangs
|
|
// up; wrote reports how many bytes it managed to send.
|
|
func endlessGateway(t *testing.T, status int, wrote *atomic.Int64) *httptest.Server {
|
|
t.Helper()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(status)
|
|
chunk := []byte(strings.Repeat("a", 4096))
|
|
for {
|
|
n, err := w.Write(chunk)
|
|
wrote.Add(int64(n))
|
|
if err != nil {
|
|
return
|
|
}
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
return srv
|
|
}
|
|
|
|
func TestClient_ASuccessBodyLargerThanTheCapIsAnErrorAndTheClientStopsReading(t *testing.T) {
|
|
var wrote atomic.Int64
|
|
srv := endlessGateway(t, http.StatusOK, &wrote)
|
|
c := New(srv.URL, "", "m", 30*time.Second, WithMaxResponseBytes(8192))
|
|
|
|
start := time.Now()
|
|
_, err := c.Complete(context.Background(), "s", "u")
|
|
|
|
if err == nil || !strings.Contains(err.Error(), "exceeds") {
|
|
t.Fatalf("want a response-too-large error, got %v", err)
|
|
}
|
|
if time.Since(start) > 5*time.Second {
|
|
t.Fatalf("the client kept reading: took %s", time.Since(start))
|
|
}
|
|
// Hanging up lets the server's writes fail; allow what the transport had already buffered.
|
|
time.Sleep(100 * time.Millisecond)
|
|
if got := wrote.Load(); got > 8192+1<<20 {
|
|
t.Fatalf("server managed to send %d bytes after a cap of 8192", got)
|
|
}
|
|
}
|
|
|
|
func TestClient_AFailingBodyIsReadOnlyUpToTheErrorBound(t *testing.T) {
|
|
var wrote atomic.Int64
|
|
srv := endlessGateway(t, http.StatusBadGateway, &wrote)
|
|
c := New(srv.URL, "", "m", 30*time.Second)
|
|
|
|
start := time.Now()
|
|
_, err := c.Complete(context.Background(), "s", "u")
|
|
|
|
he, ok := err.(*HTTPError)
|
|
if !ok {
|
|
t.Fatalf("want *HTTPError, got %T: %v", err, err)
|
|
}
|
|
if he.Code != http.StatusBadGateway {
|
|
t.Fatalf("code %d", he.Code)
|
|
}
|
|
if len(he.Body) > maxErrorBody+1 {
|
|
t.Fatalf("kept %d bytes of an error body, bound is %d", len(he.Body), maxErrorBody)
|
|
}
|
|
if time.Since(start) > 5*time.Second {
|
|
t.Fatalf("the client kept reading an endless error body: took %s", time.Since(start))
|
|
}
|
|
}
|