Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
3.2 kB · 99 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100package helps
import ( "strings" "testing"
"github.com/tidwall/gjson")
var leadingGeminiUserContentOutput []byte
func TestEnsureGeminiLeadingUserContentReusesLargeValidPayload(t *testing.T) { input := []byte(`{"contents":[{"role":"user","parts":[{"inlineData":{"mimeType":"video/mp4","data":"` + strings.Repeat("A", 4<<20) + `"}}]}]}`)
output := EnsureGeminiLeadingUserContent(input, "contents") if &output[0] != &input[0] { t.Fatal("valid request should reuse the input payload") }
result := testing.Benchmark(func(b *testing.B) { for b.Loop() { leadingGeminiUserContentOutput = EnsureGeminiLeadingUserContent(input, "contents") } }) if allocated := result.AllocedBytesPerOp(); allocated >= 1<<20 { t.Fatalf("valid 4 MiB request allocated %d bytes/op, want less than 1 MiB", allocated) }}
func TestEnsureGeminiLeadingUserContent(t *testing.T) { tests := []struct { name string inputJSON string path string wantRoles string wantLeadingEmpty bool }{ { name: "user first is unchanged", inputJSON: `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`, path: "contents", wantRoles: "user", }, { name: "leading model functionCall gets empty user", inputJSON: `{"contents":[{"role":"model","parts":[{"functionCall":{"name":"run"}}]},{"role":"user","parts":[{"functionResponse":{"name":"run"}}]}]}`, path: "contents", wantRoles: "user,model,user", wantLeadingEmpty: true, }, { name: "leading model text gets empty user and preserves following turns", inputJSON: `{"contents":[{"role":"model","parts":[{"text":"answer"}]},{"role":"user","parts":[{"text":"continue"}]}]}`, path: "contents", wantRoles: "user,model,user", wantLeadingEmpty: true, }, { name: "nested contents are normalized", inputJSON: `{"request":{"contents":[{"role":"model","parts":[{"text":"answer"}]},{"role":"user","parts":[{"text":"continue"}]}]}}`, path: "request.contents", wantRoles: "request.user,model,user", wantLeadingEmpty: true, }, { name: "empty contents are unchanged", inputJSON: `{"contents":[]}`, path: "contents", wantRoles: "", }, { name: "missing contents are unchanged", inputJSON: `{"model":"test"}`, path: "contents", wantRoles: "", }, }
for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { out := EnsureGeminiLeadingUserContent([]byte(tt.inputJSON), tt.path) contents := gjson.GetBytes(out, tt.path).Array() roles := make([]string, 0, len(contents)) for _, content := range contents { roles = append(roles, content.Get("role").String()) } expectedRoles := strings.TrimPrefix(tt.wantRoles, "request.") if got := strings.Join(roles, ","); got != expectedRoles { t.Fatalf("roles = %q, want %q; output=%s", got, expectedRoles, out) } if tt.wantLeadingEmpty { text := gjson.GetBytes(out, tt.path+".0.parts.0.text") if !text.Exists() || text.String() != "" { t.Fatalf("leading empty user part missing; output=%s", out) } } }) }}