ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/http_test.go

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