116 lines
3.9 KiB
Go
116 lines
3.9 KiB
Go
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)
|
|
}
|
|
})
|
|
}
|
|
}
|