test(rust): cover Bedrock bridge dispatch

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-29 04:05:14 +00:00
parent 30e810abc5
commit f5362ac1f0

View file

@ -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()