summaryrefslogtreecommitdiff
path: root/internal/usage/observer_test.go
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--internal/usage/observer_test.go71
1 files changed, 71 insertions, 0 deletions
diff --git a/internal/usage/observer_test.go b/internal/usage/observer_test.go
index 4a4eba6..e96a38a 100644
--- a/internal/usage/observer_test.go
+++ b/internal/usage/observer_test.go
@@ -2,6 +2,7 @@ package usage
import (
"testing"
+ "time"
"aigw/internal/domain"
)
@@ -57,3 +58,73 @@ func TestObserverReadsResponsesUsage(t *testing.T) {
t.Fatalf("unexpected streaming Responses usage: %+v reported=%v", got, stream.Reported())
}
}
+
+func TestObserverReadsEmbeddingsUsage(t *testing.T) {
+ observer := NewObserver(domain.ProtocolOpenAIEmbeddings, false)
+ _, _ = observer.Write([]byte(`{"object":"list","data":[],"usage":{"prompt_tokens":17,"total_tokens":17}}`))
+ got := observer.Usage()
+ if !observer.Reported() || got.InputTokens != 17 || got.OutputTokens != 0 || got.TotalTokens != 17 {
+ t.Fatalf("unexpected Embeddings usage: %+v reported=%v", got, observer.Reported())
+ }
+}
+
+func TestObserverMarksFirstVisibleStreamingOutput(t *testing.T) {
+ base := time.Date(2026, 8, 6, 1, 2, 3, 0, time.UTC)
+ tests := []struct {
+ name string
+ protocol domain.Protocol
+ metadata string
+ output string
+ }{
+ {
+ name: "openai chat", protocol: domain.ProtocolOpenAI,
+ metadata: "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":null}}]}\n\n",
+ output: "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n",
+ },
+ {
+ name: "responses", protocol: domain.ProtocolOpenAIResponses,
+ metadata: "data: {\"type\":\"response.created\"}\n\n",
+ output: "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello\"}\n\n",
+ },
+ {
+ name: "anthropic", protocol: domain.ProtocolAnthropic,
+ metadata: "data: {\"type\":\"content_block_start\",\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
+ output: "data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n",
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ observer := NewObserver(test.protocol, true)
+ current := base
+ observer.now = func() time.Time { return current }
+ _, _ = observer.Write([]byte(test.metadata))
+ if got := observer.FirstOutputAt(); !got.IsZero() {
+ t.Fatalf("metadata marked first output at %v", got)
+ }
+ current = base.Add(275 * time.Millisecond)
+ _, _ = observer.Write([]byte(test.output))
+ if got := observer.FirstOutputAt(); !got.Equal(current) {
+ t.Fatalf("first output = %v, want %v", got, current)
+ }
+ current = base.Add(time.Second)
+ _, _ = observer.Write([]byte(test.output))
+ if got := observer.FirstOutputAt(); !got.Equal(base.Add(275 * time.Millisecond)) {
+ t.Fatalf("first output changed to %v", got)
+ }
+ })
+ }
+}
+
+func TestObserverMarksFirstNonStreamingBodyWrite(t *testing.T) {
+ base := time.Date(2026, 8, 6, 1, 2, 3, 0, time.UTC)
+ observer := NewObserver(domain.ProtocolOpenAI, false)
+ observer.now = func() time.Time { return base }
+ _, _ = observer.Write(nil)
+ if !observer.FirstOutputAt().IsZero() {
+ t.Fatal("empty write must not mark first output")
+ }
+ _, _ = observer.Write([]byte(`{"choices":[]}`))
+ if got := observer.FirstOutputAt(); !got.Equal(base) {
+ t.Fatalf("first output = %v, want %v", got, base)
+ }
+}