internal/adapters/llm/bodycap_test.go

@@ -0,0 +1,81 @@
+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))
+	}
+}