mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
test(rust): cover Bedrock bridge dispatch
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
parent
30e810abc5
commit
f5362ac1f0
1 changed files with 27 additions and 1 deletions
|
|
@ -213,12 +213,13 @@ async def test_amessages_wrapper_forwards_args():
|
|||
custom_llm_provider="azure_ai",
|
||||
extra_headers=None,
|
||||
timeout=12.5,
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
assert response == FAKE_MESSAGES_RESPONSE
|
||||
assert bridge.calls[0]["model"] == "claude-sonnet-4-5"
|
||||
assert bridge.calls[0]["timeout_seconds"] == 12.5
|
||||
assert bridge.calls[0]["aws_region_name"] is None
|
||||
assert bridge.calls[0]["aws_region_name"] == "us-east-1"
|
||||
|
||||
|
||||
def _gate(**overrides):
|
||||
|
|
@ -257,6 +258,31 @@ async def test_gate_invokes_rust_and_marks_response_header():
|
|||
assert call["aws_region_name"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_invokes_rust_for_bedrock_and_forwards_region():
|
||||
bridge = RecordingAsyncMessages()
|
||||
litellm.use_litellm_rust(True, amessages=bridge)
|
||||
|
||||
response = await _gate(
|
||||
custom_llm_provider="bedrock",
|
||||
litellm_params=GenericLiteLLMParams(
|
||||
api_key="bedrock-token",
|
||||
aws_region_name="us-east-1",
|
||||
rust=True,
|
||||
),
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
api_key="bedrock-token",
|
||||
api_base=None,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert bridge.calls[0]["model"] == "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
assert bridge.calls[0]["custom_llm_provider"] == "bedrock"
|
||||
assert bridge.calls[0]["aws_region_name"] == "us-east-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_falls_back_to_python_when_bridge_raises():
|
||||
bridge = RaisingAsyncMessages()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue