mirror of
https://github.com/danielmiessler/Fabric.git
synced 2026-01-08 22:08:03 -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)
|
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