wangbo
add99e5421
仅对已认证 Acceptance Run 强制下载并持久化 URL 型结果，确保集群外压测端能验证最终视频；普通生产任务继续遵循现有文件存储策略。压测器仅对 Gateway 自有地址做双入口重写，避免向第三方媒体域名泄露验收凭据。\n\n同时修复本地集群重复使用同一快照文件时的幂等复制失败。验证：Go 全量测试、go vet、gofmt、bash -n、ShellCheck 和 git diff --check 通过。
241 lines
8.1 KiB
Go
241 lines
8.1 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"sync/atomic"
|
|
"testing"
|
|
)
|
|
|
|
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 TestStreamGeminiRequestBodyPreservesInput(t *testing.T) {
|
|
input := paddedPNGVariant(2<<20, 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(input)
|
|
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) != 2 || body.Contents[0].Parts[1].InlineData == nil {
|
|
t.Fatalf("unexpected Gemini body structure: %+v", body)
|
|
}
|
|
inlineData := body.Contents[0].Parts[1].InlineData
|
|
if inlineData.MIMEType != "image/png" {
|
|
t.Fatalf("mime type=%q", inlineData.MIMEType)
|
|
}
|
|
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
|
|
if err != nil {
|
|
t.Fatalf("decode streamed input: %v", err)
|
|
}
|
|
if !bytes.Equal(decoded, input) {
|
|
t.Fatal("streamed input differs from source")
|
|
}
|
|
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
|
|
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
|
|
}
|
|
}
|
|
|
|
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))
|
|
}
|
|
}
|
|
|
|
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, 256<<10, 256<<10, false)
|
|
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 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 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 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")
|
|
}
|
|
}
|