test(bedrock): add unit tests for ChatStream/parseStreamResponse
Tests cover: text-only streaming with chunk accumulation, tool call parsing with fragmented JSON, mixed text+tool responses, context cancellation, invalid JSON fallback to raw payload, nil stream guard, default finish reason, and all stop reason mappings. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Andy Lo-A-Foe
parent
ad5232ade8
commit
b03fa61764
1 changed files with 270 additions and 0 deletions
|
|
@ -8,6 +8,7 @@
|
|||
package bedrock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
|
|
@ -605,3 +606,272 @@ func TestIsSSOTokenError(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
// mockStreamReader implements bedrockruntime.ConverseStreamOutputReader for testing.
|
||||
type mockStreamReader struct {
|
||||
ch chan types.ConverseStreamOutput
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *mockStreamReader) Events() <-chan types.ConverseStreamOutput { return r.ch }
|
||||
func (r *mockStreamReader) Close() error { return nil }
|
||||
func (r *mockStreamReader) Err() error { return r.err }
|
||||
|
||||
func newMockStream(events []types.ConverseStreamOutput) *bedrockruntime.ConverseStreamEventStream {
|
||||
ch := make(chan types.ConverseStreamOutput, len(events))
|
||||
for _, e := range events {
|
||||
ch <- e
|
||||
}
|
||||
close(ch)
|
||||
|
||||
return bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
|
||||
es.Reader = &mockStreamReader{ch: ch}
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_TextOnly(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "Hello "},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "World"},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonEndTurn},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMetadata{
|
||||
Value: types.ConverseStreamMetadataEvent{
|
||||
Usage: &types.TokenUsage{
|
||||
InputTokens: aws.Int32(10),
|
||||
OutputTokens: aws.Int32(5),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
var chunks []string
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
|
||||
chunks = append(chunks, accumulated)
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Hello World", resp.Content)
|
||||
assert.Equal(t, "stop", resp.FinishReason)
|
||||
assert.Empty(t, resp.ToolCalls)
|
||||
require.NotNil(t, resp.Usage)
|
||||
assert.Equal(t, 10, resp.Usage.PromptTokens)
|
||||
assert.Equal(t, 5, resp.Usage.CompletionTokens)
|
||||
assert.Equal(t, 15, resp.Usage.TotalTokens)
|
||||
assert.Equal(t, []string{"Hello ", "Hello World"}, chunks)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_ToolCall(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_1"),
|
||||
Name: aws.String("search"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{"q":`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`"test"}`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "tool_calls", resp.FinishReason)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, "call_1", resp.ToolCalls[0].ID)
|
||||
assert.Equal(t, "search", resp.ToolCalls[0].Name)
|
||||
assert.Equal(t, map[string]any{"q": "test"}, resp.ToolCalls[0].Arguments)
|
||||
require.NotNil(t, resp.ToolCalls[0].Function)
|
||||
assert.Equal(t, "search", resp.ToolCalls[0].Function.Name)
|
||||
assert.Equal(t, `{"q":"test"}`, resp.ToolCalls[0].Function.Arguments)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_TextAndToolCall(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "Let me search that."},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(1),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_2"),
|
||||
Name: aws.String("web"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(1),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{"url":"https://example.com"}`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(1)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
var chunks []string
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
|
||||
chunks = append(chunks, accumulated)
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Let me search that.", resp.Content)
|
||||
assert.Equal(t, "tool_calls", resp.FinishReason)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, "web", resp.ToolCalls[0].Name)
|
||||
assert.Equal(t, []string{"Let me search that."}, chunks)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_ContextCancelled(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
// Use an unbuffered channel with no events so ctx.Done() is the only ready case.
|
||||
ch := make(chan types.ConverseStreamOutput)
|
||||
|
||||
stream := bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
|
||||
es.Reader = &mockStreamReader{ch: ch}
|
||||
})
|
||||
|
||||
_, err := parseStreamResponse(ctx, stream, nil)
|
||||
assert.ErrorIs(t, err, context.Canceled)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_InvalidToolJSON(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockStart{
|
||||
Value: types.ContentBlockStartEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Start: &types.ContentBlockStartMemberToolUse{
|
||||
Value: types.ToolUseBlockStart{
|
||||
ToolUseId: aws.String("call_bad"),
|
||||
Name: aws.String("broken"),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
Delta: &types.ContentBlockDeltaMemberToolUse{
|
||||
Value: types.ToolUseBlockDelta{Input: aws.String(`{not valid json`)},
|
||||
},
|
||||
},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberContentBlockStop{
|
||||
Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
|
||||
},
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, resp.ToolCalls, 1)
|
||||
assert.Equal(t, map[string]any{"raw": `{not valid json`}, resp.ToolCalls[0].Arguments)
|
||||
assert.JSONEq(t, `{"raw":"{not valid json"}`, resp.ToolCalls[0].Function.Arguments)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_DefaultFinishReason(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberContentBlockDelta{
|
||||
Value: types.ContentBlockDeltaEvent{
|
||||
Delta: &types.ContentBlockDeltaMemberText{Value: "partial"},
|
||||
ContentBlockIndex: aws.Int32(0),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "stop", resp.FinishReason)
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_NilStream(t *testing.T) {
|
||||
_, err := parseStreamResponse(context.Background(), nil, nil)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "nil event stream")
|
||||
}
|
||||
|
||||
func TestParseStreamResponse_StopReasons(t *testing.T) {
|
||||
tests := []struct {
|
||||
reason types.StopReason
|
||||
expected string
|
||||
}{
|
||||
{types.StopReasonEndTurn, "stop"},
|
||||
{types.StopReasonMaxTokens, "length"},
|
||||
{types.StopReasonToolUse, "tool_calls"},
|
||||
{types.StopReasonStopSequence, "stop"},
|
||||
{types.StopReasonContentFiltered, "content_filter"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(string(tt.reason), func(t *testing.T) {
|
||||
events := []types.ConverseStreamOutput{
|
||||
&types.ConverseStreamOutputMemberMessageStop{
|
||||
Value: types.MessageStopEvent{StopReason: tt.reason},
|
||||
},
|
||||
}
|
||||
stream := newMockStream(events)
|
||||
resp, err := parseStreamResponse(context.Background(), stream, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tt.expected, resp.FinishReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue