test: add test for metadata stream propagation

This commit is contained in:
berniegreen
2025-12-31 15:56:20 -06:00
parent 07457d86d3
commit c7ecac3262

View File

@@ -232,3 +232,83 @@ func TestChatter_Send_StreamingSuccessfulAggregation(t *testing.T) {
t.Errorf("Expected aggregated message %q, got %q", expectedMessage, assistantMessage.Content)
}
}
func TestChatter_Send_StreamingMetadataPropagation(t *testing.T) {
// Create a temporary database for testing
tempDir := t.TempDir()
db := fsdb.NewDb(tempDir)
// Create test chunks: one content, one usage metadata
testChunks := []domain.StreamUpdate{
{
Type: domain.StreamTypeContent,
Content: "Test content",
},
{
Type: domain.StreamTypeUsage,
Usage: &domain.UsageMetadata{
InputTokens: 10,
OutputTokens: 5,
TotalTokens: 15,
},
},
}
// Create a mock vendor
mockVendor := &mockVendor{
sendStreamError: nil,
streamChunks: testChunks,
}
// Create chatter with streaming enabled
chatter := &Chatter{
db: db,
Stream: true,
vendor: mockVendor,
model: "test-model",
}
// Create a test request
request := &domain.ChatRequest{
Message: &chat.ChatCompletionMessage{
Role: chat.ChatMessageRoleUser,
Content: "test message",
},
}
// Create an update channel to capture stream events
updateChan := make(chan domain.StreamUpdate, 10)
// Create test options with UpdateChan
opts := &domain.ChatOptions{
Model: "test-model",
UpdateChan: updateChan,
Quiet: true, // Suppress stdout/stderr
}
// Call Send
_, err := chatter.Send(request, opts)
if err != nil {
t.Fatalf("Expected no error, but got: %v", err)
}
close(updateChan)
// Verify we received the metadata event
var usageReceived bool
for update := range updateChan {
if update.Type == domain.StreamTypeUsage {
usageReceived = true
if update.Usage == nil {
t.Error("Expected usage metadata to be non-nil")
} else {
if update.Usage.TotalTokens != 15 {
t.Errorf("Expected 15 total tokens, got %d", update.Usage.TotalTokens)
}
}
}
}
if !usageReceived {
t.Error("Expected to receive a usage metadata update, but didn't")
}
}