90 lines
3.0 KiB
Go
90 lines
3.0 KiB
Go
package protocol
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestMergeUsageSizeBoundary(t *testing.T) {
|
|
for _, kind := range []Kind{Chat, Responses, Anthropic} {
|
|
t.Run(string(kind), func(t *testing.T) {
|
|
m := object{"padding": strings.Repeat("x", maxStreamUsageBytes-len(`{"padding":""}`))}
|
|
u, err := mergeUsage(kind, Usage{}, object{"usage": m})
|
|
if err != nil || len(u.Raw) != maxStreamUsageBytes {
|
|
t.Fatalf("exact limit: size=%d err=%v", len(u.Raw), err)
|
|
}
|
|
for _, update := range []object{
|
|
{"usage": object{"padding": str(m["padding"]) + "x"}},
|
|
{"usage": object{"new_key": 1}},
|
|
} {
|
|
got, err := mergeUsage(kind, u, update)
|
|
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") || !reflect.DeepEqual(got, u) {
|
|
t.Fatalf("overflow: retained=%d err=%v", len(got.Raw), err)
|
|
}
|
|
}
|
|
for i := 0; i < 3; i++ {
|
|
got, err := mergeUsage(kind, u, object{"usage": m})
|
|
if err != nil || !reflect.DeepEqual(got, u) {
|
|
t.Fatalf("replacement at limit: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestStreamUsageSizeLimits(t *testing.T) {
|
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
|
for _, cumulative := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("%s/%s/cumulative=%t", from, to, cumulative), func(t *testing.T) {
|
|
var input strings.Builder
|
|
writeUsage := func(u object) {
|
|
m := object{"usage": u}
|
|
switch from {
|
|
case Anthropic:
|
|
m["type"] = "ping"
|
|
case Responses:
|
|
m = object{"type": "response.in_progress", "response": m}
|
|
}
|
|
input.WriteString(limitEvent(mustLimitJSON(t, m)))
|
|
}
|
|
writeUsage(object{"prompt_tokens": 7, "input_tokens": 7})
|
|
if cumulative {
|
|
for i := 0; i < 32; i++ {
|
|
writeUsage(object{fmt.Sprintf("key_%d", i): strings.Repeat("x", 1024)})
|
|
}
|
|
} else {
|
|
writeUsage(object{"padding": strings.Repeat("x", maxStreamUsageBytes)})
|
|
}
|
|
u, err := Stream(context.Background(), io.Discard, strings.NewReader(input.String()), from, to, "")
|
|
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") {
|
|
t.Fatalf("expected usage limit error, got %v", err)
|
|
}
|
|
if !u.Known || u.Input != 7 || len(u.Raw) > maxStreamUsageBytes {
|
|
t.Fatalf("lost valid usage: known=%t input=%d size=%d", u.Known, u.Input, len(u.Raw))
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestMergeUsagePartialUpdates(t *testing.T) {
|
|
u, err := mergeUsage(Anthropic, Usage{}, object{"message": object{"usage": object{"input_tokens": 5, "cache_read_input_tokens": 2}}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
u, err = mergeUsage(Anthropic, u, object{"usage": object{"output_tokens": 3}})
|
|
if err != nil || !u.Known || u.Input != 7 || u.Output != 3 || u.CacheRead != 2 {
|
|
t.Fatalf("partial update: %+v err=%v", u, err)
|
|
}
|
|
got, err := mergeUsage(Anthropic, u, object{})
|
|
if err != nil || !reflect.DeepEqual(got, u) {
|
|
t.Fatalf("missing usage changed totals: %+v err=%v", got, err)
|
|
}
|
|
}
|