mirror of
https://github.com/danielmiessler/Fabric.git
synced 2026-01-07 21:44:02 -05:00
test: add test for metadata stream propagation
This commit is contained in:
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user