wangbo
c28bf74230
保留平台模型 RPM、TPM 和并发策略语义，增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载，避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标，并扩展本地 acceptance runner，覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证：go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过；48/48 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
407 lines
14 KiB
Go
407 lines
14 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"sync/atomic"
|
|
"testing"
|
|
)
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
|
return f(request)
|
|
}
|
|
|
|
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
|
|
payload := paddedPNG(4 << 20)
|
|
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
|
|
base64.StdEncoding.EncodeToString(payload))
|
|
size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
|
|
if err != nil {
|
|
t.Fatalf("stream Gemini output: %v", err)
|
|
}
|
|
expected := sha256.Sum256(payload)
|
|
if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
|
|
t.Fatalf("size=%d digest=%s", size, digest)
|
|
}
|
|
}
|
|
|
|
func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
|
|
first := paddedPNGVariant(256<<10, 1)
|
|
second := paddedPNGVariant(256<<10, 2)
|
|
if len(first) != 256<<10 || len(second) != 256<<10 {
|
|
t.Fatalf("variant sizes=%d/%d", len(first), len(second))
|
|
}
|
|
firstHash := sha256.Sum256(first)
|
|
secondHash := sha256.Sum256(second)
|
|
if firstHash == secondHash {
|
|
t.Fatal("distinct variants have the same SHA-256")
|
|
}
|
|
}
|
|
|
|
func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
|
|
inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
|
|
var body struct {
|
|
Contents []struct {
|
|
Parts []struct {
|
|
InlineData *struct {
|
|
MIMEType string `json:"mimeType"`
|
|
Data string `json:"data"`
|
|
} `json:"inlineData"`
|
|
} `json:"parts"`
|
|
} `json:"contents"`
|
|
GenerationConfig struct {
|
|
ResponseModalities []string `json:"responseModalities"`
|
|
} `json:"generationConfig"`
|
|
}
|
|
stream := streamGeminiRequestBody(inputs)
|
|
defer stream.Close()
|
|
if err := json.NewDecoder(stream).Decode(&body); err != nil {
|
|
t.Fatalf("decode streamed Gemini request: %v", err)
|
|
}
|
|
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
|
|
t.Fatalf("unexpected Gemini body structure: %+v", body)
|
|
}
|
|
for index, input := range inputs {
|
|
inlineData := body.Contents[0].Parts[index+1].InlineData
|
|
if inlineData == nil || inlineData.MIMEType != "image/png" {
|
|
t.Fatalf("image %d inline data=%+v", index, inlineData)
|
|
}
|
|
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
|
|
if err != nil {
|
|
t.Fatalf("decode streamed input %d: %v", index, err)
|
|
}
|
|
if !bytes.Equal(decoded, input) {
|
|
t.Fatalf("streamed input %d differs from source", index)
|
|
}
|
|
}
|
|
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
|
|
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
|
|
}
|
|
}
|
|
|
|
func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
|
|
inputs := geminiInputs(3, 768<<10, "replay-test", 91)
|
|
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
|
|
if err != nil {
|
|
t.Fatalf("new Gemini request: %v", err)
|
|
}
|
|
if request.GetBody == nil {
|
|
t.Fatal("streaming Gemini request does not provide GetBody")
|
|
}
|
|
first, err := io.ReadAll(request.Body)
|
|
if err != nil {
|
|
t.Fatalf("read first request body: %v", err)
|
|
}
|
|
_ = request.Body.Close()
|
|
replayedBody, err := request.GetBody()
|
|
if err != nil {
|
|
t.Fatalf("replay Gemini request body: %v", err)
|
|
}
|
|
replayed, err := io.ReadAll(replayedBody)
|
|
if err != nil {
|
|
t.Fatalf("read replayed request body: %v", err)
|
|
}
|
|
_ = replayedBody.Close()
|
|
if !bytes.Equal(first, replayed) {
|
|
t.Fatal("replayed Gemini request body differs from the original")
|
|
}
|
|
if request.ContentLength != int64(len(first)) {
|
|
t.Fatalf("content length=%d, want %d", request.ContentLength, len(first))
|
|
}
|
|
}
|
|
|
|
func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
|
|
request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
|
|
"https://fixtures.example/one.png",
|
|
"https://fixtures.example/two.png",
|
|
"https://fixtures.example/three.png",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("new multi-image request: %v", err)
|
|
}
|
|
var body struct {
|
|
Contents []struct {
|
|
Parts []struct {
|
|
FileData *struct {
|
|
FileURI string `json:"fileUri"`
|
|
} `json:"fileData"`
|
|
} `json:"parts"`
|
|
} `json:"contents"`
|
|
}
|
|
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
|
|
t.Fatalf("decode multi-image request: %v", err)
|
|
}
|
|
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
|
|
t.Fatalf("unexpected multi-image body: %+v", body)
|
|
}
|
|
for index := 1; index < 4; index++ {
|
|
if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
|
|
t.Fatalf("missing fileData at part %d", index)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
|
|
images := make([]string, 16)
|
|
for index := range images {
|
|
images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
|
|
}
|
|
combinations := videoCombinations(images, 128)
|
|
if len(combinations) != 128 {
|
|
t.Fatalf("combinations=%d", len(combinations))
|
|
}
|
|
seen := map[string]struct{}{}
|
|
for _, combination := range combinations {
|
|
seen[fmt.Sprint(combination)] = struct{}{}
|
|
if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
|
|
t.Fatalf("invalid combination size=%d", len(combination))
|
|
}
|
|
}
|
|
if len(seen) != 128 {
|
|
t.Fatalf("unique combinations=%d", len(seen))
|
|
}
|
|
if got := requestedVideoCombinationCount(combinations, 144, 0, 1); got != 131 {
|
|
t.Fatalf("requested combinations=%d, want 131", got)
|
|
}
|
|
}
|
|
|
|
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
|
|
output := paddedPNG(256 << 10)
|
|
encoded := base64.StdEncoding.EncodeToString(output)
|
|
var firstCalls atomic.Int64
|
|
var secondCalls atomic.Int64
|
|
var firstKeyCalls atomic.Int64
|
|
var secondKeyCalls atomic.Int64
|
|
newGateway := func(calls *atomic.Int64) *httptest.Server {
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls.Add(1)
|
|
switch r.Header.Get("Authorization") {
|
|
case "Bearer key-1":
|
|
firstKeyCalls.Add(1)
|
|
case "Bearer key-2":
|
|
secondKeyCalls.Add(1)
|
|
default:
|
|
t.Errorf("unexpected authorization header")
|
|
}
|
|
if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
|
|
t.Errorf("missing acceptance headers")
|
|
}
|
|
var body map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Errorf("decode Gemini body: %v", err)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"candidates": []any{map[string]any{"content": map[string]any{
|
|
"parts": []any{map[string]any{"inlineData": map[string]any{
|
|
"mimeType": "image/png", "data": encoded,
|
|
}}},
|
|
}}},
|
|
})
|
|
}))
|
|
}
|
|
first := newGateway(&firstCalls)
|
|
defer first.Close()
|
|
second := newGateway(&secondCalls)
|
|
defer second.Close()
|
|
|
|
opts := options{
|
|
gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
|
|
runToken: "token-1", geminiModel: "gemini-image-test",
|
|
}
|
|
result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1)
|
|
if result.err != nil || result.report.Completed != 8 {
|
|
t.Fatalf("result=%+v err=%v", result.report, result.err)
|
|
}
|
|
if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
|
|
t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
|
|
}
|
|
if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
|
|
t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
|
|
}
|
|
}
|
|
|
|
func TestDistributedShardIndexesCoverWorkloadWithoutOverlap(t *testing.T) {
|
|
first := options{shardIndex: 0, shardCount: 2}
|
|
second := options{shardIndex: 1, shardCount: 2}
|
|
if first.shardRequestCount(1001) != 501 || second.shardRequestCount(1001) != 500 {
|
|
t.Fatalf("shard counts=%d/%d", first.shardRequestCount(1001), second.shardRequestCount(1001))
|
|
}
|
|
seen := map[int]bool{}
|
|
for local := 0; local < first.shardRequestCount(1001); local++ {
|
|
seen[first.logicalRequestIndex(local)] = true
|
|
}
|
|
for local := 0; local < second.shardRequestCount(1001); local++ {
|
|
index := second.logicalRequestIndex(local)
|
|
if seen[index] {
|
|
t.Fatalf("duplicate logical index %d", index)
|
|
}
|
|
seen[index] = true
|
|
}
|
|
if len(seen) != 1001 {
|
|
t.Fatalf("covered indexes=%d", len(seen))
|
|
}
|
|
}
|
|
|
|
func TestGeminiInputVariantIsRunScoped(t *testing.T) {
|
|
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-b", 7) {
|
|
t.Fatal("different Run IDs produced the same input variant")
|
|
}
|
|
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-a", 8) {
|
|
t.Fatal("different logical indexes produced the same input variant")
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceIdempotencyKeyIncludesExecutionAndLogicalIndex(t *testing.T) {
|
|
request := httptest.NewRequest(http.MethodPost, "https://gateway.example/v1", nil)
|
|
opts := options{
|
|
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1", executionID: "p24-2-gemini-baseline",
|
|
}
|
|
opts.setHeaders(request, 17, false)
|
|
if got := request.Header.Get("Idempotency-Key"); got != "acceptance-run-1-p24-2-gemini-baseline-17" {
|
|
t.Fatalf("idempotency key=%q", got)
|
|
}
|
|
poll := httptest.NewRequest(http.MethodGet, "https://gateway.example/result", nil)
|
|
opts.setHeaders(poll, 17, false)
|
|
if got := poll.Header.Get("Idempotency-Key"); got != "" {
|
|
t.Fatalf("GET request unexpectedly has idempotency key %q", got)
|
|
}
|
|
}
|
|
|
|
func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
|
if request.Header.Get("Authorization") != "" {
|
|
t.Fatal("acceptance credentials leaked to external media host")
|
|
}
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
|
}))
|
|
defer server.Close()
|
|
if err := validateVideoAsset(t.Context(), server.Client(), options{}, server.URL+"/result.mp4", 0); err != nil {
|
|
t.Fatalf("validate video asset: %v", err)
|
|
}
|
|
if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
|
|
t.Fatalf("media URL=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
|
}))
|
|
defer server.Close()
|
|
poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return nil, errors.New("gateway-only transport must not be used for external media")
|
|
})}
|
|
opts := options{gatewayTLSName: "gateway.easyai.local"}
|
|
if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
|
|
t.Fatalf("validate external video with independent transport: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
|
if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {
|
|
t.Fatalf("host=%q authorization=%q", request.Host, request.Header.Get("Authorization"))
|
|
}
|
|
if request.URL.Path != "/static/generated/result.mp4" {
|
|
t.Fatalf("path=%q", request.URL.Path)
|
|
}
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
|
}))
|
|
defer server.Close()
|
|
opts := options{
|
|
gateways: []string{server.URL}, apiKeys: []string{"key-1"},
|
|
runID: "run-1", runToken: "token-1", gatewayTLSName: "gateway.easyai.local",
|
|
}
|
|
if err := validateVideoAsset(t.Context(), server.Client(), opts, "/static/generated/result.mp4", 0); err != nil {
|
|
t.Fatalf("validate materialized video: %v", err)
|
|
}
|
|
got := findMediaURL(map[string]any{
|
|
"raw": map[string]any{"video_url": "http://internal.invalid/video.mp4"},
|
|
"data": []any{map[string]any{"video_url": "/static/generated/result.mp4"}},
|
|
})
|
|
if got != "/static/generated/result.mp4" {
|
|
t.Fatalf("preferred media URL=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
|
|
got := redactError(
|
|
`token-1 failed at https://example.invalid/video.mp4?token=signed`,
|
|
options{runToken: "token-1"},
|
|
)
|
|
if got != `[REDACTED] failed at [REDACTED_URL]` {
|
|
t.Fatalf("redacted error=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceReportCannotOverwriteExistingRunArtifact(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "load-report.json")
|
|
if err := writeReportExclusive(path, []byte(`{"runId":"first"}`)); err != nil {
|
|
t.Fatalf("write first report: %v", err)
|
|
}
|
|
if err := writeReportExclusive(path, []byte(`{"runId":"second"}`)); err == nil {
|
|
t.Fatal("second report unexpectedly overwrote the first report")
|
|
}
|
|
payload, err := os.ReadFile(path)
|
|
if err != nil {
|
|
t.Fatalf("read report: %v", err)
|
|
}
|
|
if string(payload) != "{\"runId\":\"first\"}\n" {
|
|
t.Fatalf("report changed after rejected overwrite: %s", payload)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceFailurePreservesOperationAndHTTPStatus(t *testing.T) {
|
|
err := withOperation("video_poll", &httpStatusError{Status: http.StatusUnauthorized, Body: "unauthorized"})
|
|
var operationErr *operationError
|
|
if !errors.As(err, &operationErr) || operationErr.Operation != "video_poll" {
|
|
t.Fatalf("operation error=%#v", operationErr)
|
|
}
|
|
var statusErr *httpStatusError
|
|
if !errors.As(err, &statusErr) || statusErr.Status != http.StatusUnauthorized {
|
|
t.Fatalf("HTTP status error=%#v", statusErr)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
|
|
request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
|
|
opts := options{
|
|
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
|
|
gatewayTLSName: "ai.example.com",
|
|
}
|
|
opts.setHeaders(request, 0, false)
|
|
if request.Host != "ai.example.com" {
|
|
t.Fatalf("Host=%q", request.Host)
|
|
}
|
|
}
|
|
|
|
func TestAcceptanceRootCAsRejectsSymlink(t *testing.T) {
|
|
root := t.TempDir()
|
|
target := filepath.Join(root, "ca.pem")
|
|
if err := os.WriteFile(target, []byte("not a certificate"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
link := filepath.Join(root, "ca-link.pem")
|
|
if err := os.Symlink(target, link); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := acceptanceRootCAs(link); err == nil {
|
|
t.Fatal("symlink CA file was accepted")
|
|
}
|
|
}
|