ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/front_test.go

@@ -0,0 +1,116 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func testFrontend(url string) frontend {
+	return frontend{client: http.DefaultClient, gateway: url, prefix: "batch/", routePrefix: "native/", heartbeat: 10 * time.Millisecond}
+}
+
+func TestFrontendPayload(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var body map[string]json.RawMessage
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Error(err)
+		}
+		if string(body["model"]) != `"native/partner-model@001"` || string(body["tools"]) != `[{"opaque":true}]` || r.Header.Get("Authorization") != "Bearer synthetic" {
+			t.Errorf("unexpected request: %s", body)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusBadGateway)
+		_, _ = io.WriteString(w, `{"error":"synthetic failure"}`)
+	}))
+	defer upstream.Close()
+	req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/partner-model@001","tools":[{"opaque":true}]}`))
+	req.Header.Set("Authorization", "Bearer synthetic")
+	res := httptest.NewRecorder()
+	testFrontend(upstream.URL).ServeHTTP(res, req)
+	if res.Code != http.StatusBadGateway || res.Body.String() != `{"error":"synthetic failure"}` {
+		t.Fatalf("unexpected response: %d %s", res.Code, res.Body)
+	}
+}
+
+func TestFrontendCancel(t *testing.T) {
+	for _, headers := range []bool{false, true} {
+		t.Run(map[bool]string{false: "pending_headers", true: "pending_body"}[headers], func(t *testing.T) {
+			entered := make(chan struct{})
+			cancelled := make(chan struct{})
+			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				_, _ = io.Copy(io.Discard, r.Body)
+				if headers {
+					w.Header().Set("Content-Type", "text/event-stream")
+					_ = http.NewResponseController(w).Flush()
+				}
+				close(entered)
+				<-r.Context().Done()
+				close(cancelled)
+			}))
+			defer upstream.Close()
+			defer upstream.CloseClientConnections()
+			server := httptest.NewServer(testFrontend(upstream.URL))
+			defer server.Close()
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
+			if err != nil {
+				t.Fatal(err)
+			}
+			res, err := http.DefaultClient.Do(req)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer func() { _ = res.Body.Close() }()
+			select {
+			case <-entered:
+			case <-ctx.Done():
+				t.Fatal("upstream was not entered")
+			}
+			data := make([]byte, 1)
+			if _, err := res.Body.Read(data); err != nil {
+				t.Fatal(err)
+			}
+			if err := res.Body.Close(); err != nil {
+				t.Fatal(err)
+			}
+			select {
+			case <-cancelled:
+			case <-ctx.Done():
+				t.Fatal("disconnect did not cancel upstream")
+			}
+		})
+	}
+}
+
+func TestFrontendStream(t *testing.T) {
+	for _, complete := range []bool{false, true} {
+		t.Run(map[bool]string{false: "incomplete", true: "complete"}[complete], func(t *testing.T) {
+			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "text/event-stream")
+				if complete {
+					_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n")
+				}
+				_, _ = io.WriteString(w, "data: [DONE]\n\n")
+			}))
+			defer upstream.Close()
+			req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
+			res := httptest.NewRecorder()
+			testFrontend(upstream.URL).ServeHTTP(res, req)
+			output := res.Body.String()
+			if complete {
+				if !strings.Contains(output, "answer") || !strings.Contains(output, "[DONE]") {
+					t.Fatal(output)
+				}
+			} else if !strings.Contains(output, "batch_error") || strings.Contains(output, "[DONE]") {
+				t.Fatal(output)
+			}
+		})
+	}
+}