Skip to content

Commit c9af528

Browse files
committed
Python: test that streaming and non-streaming report the same Bedrock model
1 parent f7c5e86 commit c9af528

1 file changed

Lines changed: 36 additions & 0 deletions

File tree

‎python/packages/bedrock/tests/test_bedrock_client.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,42 @@ async def test_get_response_reports_the_requested_model(client_model: str | None
134134
assert response.model == "us.openai.gpt-6-sol"
135135

136136

137+
class _StubConverseAndStreamWithoutModelId(_StubConverseWithoutModelId):
138+
def converse_stream(self, **kwargs: Any) -> dict[str, Any]:
139+
self.calls.append(kwargs)
140+
return {
141+
"stream": _StubEventStream([
142+
{"messageStart": {"role": "assistant"}},
143+
{"contentBlockDelta": {"delta": {"text": "hi"}, "contentBlockIndex": 0}},
144+
{"contentBlockStop": {"contentBlockIndex": 0}},
145+
{"messageStop": {"stopReason": "end_turn"}},
146+
])
147+
}
148+
149+
150+
@pytest.mark.parametrize("stream", [False, True])
151+
async def test_streaming_and_non_streaming_report_the_same_model(stream: bool) -> None:
152+
"""Both paths must report the per-call model override, not the client default."""
153+
stub = _StubConverseAndStreamWithoutModelId()
154+
client = BedrockChatClient(
155+
model="amazon.titan-text",
156+
region="us-west-2",
157+
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
158+
)
159+
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
160+
options = {"model": "us.openai.gpt-6-sol"}
161+
162+
if stream:
163+
response_stream = client.get_response(messages=messages, options=options, stream=True)
164+
response = await response_stream.get_final_response()
165+
else:
166+
response = await client.get_response(messages=messages, options=options)
167+
168+
assert stub.calls[0]["modelId"] == "us.openai.gpt-6-sol"
169+
assert response.text == "hi"
170+
assert response.model == "us.openai.gpt-6-sol"
171+
172+
137173
async def test_stream_yields_updates_as_converse_stream_events_arrive() -> None:
138174
"""stream=True should use ConverseStream and yield text, tool call, finish and usage updates per event."""
139175
stub = _StubBedrockStreamRuntime([

0 commit comments

Comments
 (0)