akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
headroom/tests/test_bedrock_region.py
Matt Haitana 64cb46e24b
fix(proxy): pass through cross-region prefixed Bedrock model IDs directly (#2330) 
## Description

When a model ID with a cross-region prefix (`au.`, `us.`, `eu.`,
`apac.`, `global.`) is sent to the Bedrock backend, `map_model_id` was
normalising it (e.g. `au.anthropic.claude-opus-4-8` → `claude-opus-4-8`)
then re-looking it up in the discovery map. If an APPLICATION inference
profile wrapping the same foundation model existed in the account, it
would be returned — routing the request to a profile the caller is not
authorised to invoke, resulting in a 403 from Bedrock even though the
system-defined profile is reachable directly.

Cross-region prefixed IDs are already fully-qualified system-defined
profile IDs; they must pass through unchanged.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/backends/litellm.py`: added early-exit in `map_model_id` —
model IDs starting with `au.`, `us.`, `eu.`, `apac.`, or `global.` are
returned as `bedrock/<model_id>` without any discovery lookup
- `tests/test_bedrock_region.py`: two new regression tests covering the
exact failure mode and all five prefix families

## Testing

- [x] Unit tests pass (`pytest`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
tests/test_bedrock_region.py ........................................... [ 95%]
..                                                                       [100%]

45 passed in 4.45s
```

## Real Behavior Proof

- Environment: headroom 0.32.0-dev, `--backend bedrock`, `--region
ap-southeast-2`, `--bedrock-profile <BEDROCK_PROFILE>`, proxy on port
8788
- Exact command / steps: Output
  ```
# Before fix — old map_model_id logic with a contaminated discovery map:
  # Input:      au.anthropic.claude-opus-4-8
  # Normalized: claude-opus-4-8
  # Resolved:   bedrock/<application-inference-profile-arn>  <-- 403

  # After fix — cross-region prefix detected, passed through directly:
  curl -s -X POST http://localhost:8788/v1/messages \
    -H "Content-Type: application/json" \
    -H "x-api-key: sk-ant-dummy" \
    -H "anthropic-version: 2023-06-01" \
-d
'{"model":"au.anthropic.claude-opus-4-8","max_tokens":64,"messages":[{"role":"user","content":"Reply
with just: fix works"}]}'
  ```
- Observed result:
`{"type":"message","role":"assistant","content":[{"type":"text","text":"fix
works"}],"model":"au.anthropic.claude-opus-4-8","stop_reason":"end_turn",...}`
— HTTP 200, routed to `bedrock/au.anthropic.claude-opus-4-8`
(system-defined profile) rather than the APPLICATION profile ARN
- Not tested: `apac.` and `global.` prefixes against a live AWS account
(covered by unit tests only)

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes

## Additional Notes

The existing `_fetch_bedrock_inference_profiles` already filters to
`typeEquals="SYSTEM_DEFINED"` so APPLICATION profiles are not added to
the discovery map during normal startup. This fix closes the remaining
gap where a caller passes a cross-region prefixed ID directly —
previously that ID was normalised before lookup, which could
accidentally match a stale or externally-injected map entry.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: JD Davis <jd@jds-macbook-air.tail2a279.ts.net>
2026-08-11 23:54:03 -05:00

650 lines
28 KiB
Python
Raw Blame History

"""Tests for Bedrock region support and fallback model mapping.
Ensures that EU, AP, and US regions all produce valid Bedrock model IDs,
and that the proxy degrades gracefully when boto3 is unavailable or the
AWS API call fails.
"""
from unittest.mock import MagicMock, patch
from tests._dotenv import importorskip_no_env_leak
importorskip_no_env_leak("litellm")
from headroom.backends.litellm import ( # noqa: E402 (must follow importorskip)
LiteLLMBackend,
_bedrock_profiles_cache,
_bedrock_region_prefix,
_build_bedrock_fallback_map,
_fetch_bedrock_inference_profiles,
_normalize_bedrock_profile_id,
)
# =============================================================================
# Region Prefix Mapping
# =============================================================================
class TestBedrockRegionPrefix:
"""Test AWS region -> inference profile prefix mapping."""
def test_us_regions(self):
assert _bedrock_region_prefix("us-east-1") == "us"
assert _bedrock_region_prefix("us-west-2") == "us"
def test_eu_regions(self):
assert _bedrock_region_prefix("eu-central-1") == "eu"
assert _bedrock_region_prefix("eu-west-1") == "eu"
assert _bedrock_region_prefix("eu-west-3") == "eu"
assert _bedrock_region_prefix("eu-north-1") == "eu"
def test_ap_regions(self):
assert _bedrock_region_prefix("ap-southeast-1") == "apac"
assert _bedrock_region_prefix("ap-northeast-1") == "apac"
def test_unknown_region_defaults_to_us(self):
assert _bedrock_region_prefix("me-south-1") == "us"
assert _bedrock_region_prefix("sa-east-1") == "us"
# =============================================================================
# Static Fallback Model Map
# =============================================================================
class TestBuildBedrockFallbackMap:
"""Test static fallback model map construction."""
def test_us_region_uses_us_prefix(self):
model_map = _build_bedrock_fallback_map("us-east-1")
assert "claude-sonnet-4-20250514" in model_map
assert model_map["claude-sonnet-4-20250514"] == (
"bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0"
)
def test_eu_region_uses_eu_prefix(self):
model_map = _build_bedrock_fallback_map("eu-central-1")
assert "claude-sonnet-4-20250514" in model_map
assert model_map["claude-sonnet-4-20250514"] == (
"bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
)
def test_ap_region_uses_apac_prefix(self):
model_map = _build_bedrock_fallback_map("ap-southeast-1")
assert "claude-sonnet-4-20250514" in model_map
assert model_map["claude-sonnet-4-20250514"] == (
"bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0"
)
def test_all_models_present(self):
model_map = _build_bedrock_fallback_map("us-east-1")
expected_models = [
"claude-opus-4-6",
"claude-sonnet-4-6",
"claude-sonnet-4-20250514",
"claude-opus-4-20250514",
"claude-3-7-sonnet-20250219",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022",
"claude-3-opus-20240229",
"claude-3-haiku-20240307",
"claude-haiku-4-5-20251001",
]
for model in expected_models:
assert model in model_map, f"Missing model: {model}"
def test_all_values_are_valid_bedrock_format(self):
"""Every value must start with 'bedrock/' and contain 'anthropic.'."""
for region in ("us-east-1", "eu-west-1", "ap-northeast-1"):
model_map = _build_bedrock_fallback_map(region)
for name, litellm_id in model_map.items():
assert litellm_id.startswith("bedrock/"), (
f"{name}: expected bedrock/ prefix, got {litellm_id}"
)
assert "anthropic." in litellm_id, (
f"{name}: expected anthropic. in id, got {litellm_id}"
)
# =============================================================================
# Fetch with Graceful Fallback
# =============================================================================
class TestFetchBedrockInferenceProfiles:
"""Test dynamic fetch with fallback on failure."""
def setup_method(self):
"""Clear the cache before each test."""
_bedrock_profiles_cache.clear()
def test_fallback_when_boto3_import_fails(self):
"""Should return static map when boto3 is not installed."""
with patch.dict("sys.modules", {"boto3": None}):
# Force reimport failure
# Temporarily break boto3 import inside the function
original_import = (
__builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__
)
def mock_import(name, *args, **kwargs):
if name == "boto3":
raise ImportError("No module named 'boto3'")
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=mock_import):
_bedrock_profiles_cache.clear()
result = _fetch_bedrock_inference_profiles("eu-central-1")
assert len(result) > 0
# Should use EU prefix
assert result["claude-sonnet-4-20250514"] == (
"bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
)
def test_fallback_when_api_call_fails(self):
"""Should return static map when list_inference_profiles raises."""
mock_boto3 = MagicMock()
mock_client = MagicMock()
mock_client.list_inference_profiles.side_effect = Exception(
"AccessDeniedException: not authorized"
)
mock_session = MagicMock()
mock_session.client.return_value = mock_client
mock_boto3.Session.return_value = mock_session
with patch("headroom.backends.litellm.boto3", mock_boto3, create=True):
# Patch the import inside the function
_fetch_bedrock_inference_profiles.__code__ # noqa: B018
_bedrock_profiles_cache.clear()
# We need to actually test the function, so let's just use the
# mock_boto3 and make sure the function catches the exception
import builtins
real_import = builtins.__import__
def patched_import(name, *args, **kwargs):
if name == "boto3":
return mock_boto3
return real_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=patched_import):
_bedrock_profiles_cache.clear()
result = _fetch_bedrock_inference_profiles("eu-west-1")
assert len(result) > 0
# Should use EU prefix
for litellm_id in result.values():
assert "eu.anthropic." in litellm_id
def test_successful_fetch_uses_api_results(self):
"""When API works, should use dynamic results (not fallback)."""
mock_boto3 = MagicMock()
mock_client = MagicMock()
mock_client.list_inference_profiles.return_value = {
"inferenceProfileSummaries": [
{"inferenceProfileId": "eu.anthropic.claude-sonnet-4-20250514-v1:0"},
{"inferenceProfileId": "eu.anthropic.claude-3-5-sonnet-20241022-v2:0"},
{"inferenceProfileId": "eu.meta.llama-3-70b-v1:0"}, # non-Anthropic, should skip
]
}
mock_session = MagicMock()
mock_session.client.return_value = mock_client
mock_boto3.Session.return_value = mock_session
import builtins
real_import = builtins.__import__
def patched_import(name, *args, **kwargs):
if name == "boto3":
return mock_boto3
return real_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=patched_import):
_bedrock_profiles_cache.clear()
result = _fetch_bedrock_inference_profiles("eu-central-1")
assert len(result) == 2
assert result["claude-sonnet-4-20250514"] == (
"bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
)
assert result["claude-3-5-sonnet-20241022"] == (
"bedrock/eu.anthropic.claude-3-5-sonnet-20241022-v2:0"
)
def test_caching_prevents_repeated_api_calls(self):
"""Second call for same region+profile should return cached result."""
_bedrock_profiles_cache.clear()
_bedrock_profiles_cache["us-east-1:"] = {"test": "bedrock/test-model"}
result = _fetch_bedrock_inference_profiles("us-east-1")
assert result == {"test": "bedrock/test-model"}
def test_profile_cache_isolation(self):
"""Different profiles for the same region must not share a cache entry."""
_bedrock_profiles_cache.clear()
_bedrock_profiles_cache["us-east-1:profileA"] = {"model": "bedrock/profile-a-model"}
_bedrock_profiles_cache["us-east-1:profileB"] = {"model": "bedrock/profile-b-model"}
result_a = _fetch_bedrock_inference_profiles("us-east-1", profile_name="profileA")
result_b = _fetch_bedrock_inference_profiles("us-east-1", profile_name="profileB")
assert result_a["model"] == "bedrock/profile-a-model"
assert result_b["model"] == "bedrock/profile-b-model"
assert result_a != result_b
# =============================================================================
# LiteLLMBackend.map_model_id with EU Regions
# =============================================================================
class TestBedrockModelMapping:
"""Test model ID mapping for different regions."""
def setup_method(self):
_bedrock_profiles_cache.clear()
def test_eu_region_maps_correctly(self):
"""EU region should produce eu.anthropic.* model IDs."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={
"claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0",
},
):
backend = LiteLLMBackend(provider="bedrock", region="eu-central-1")
result = backend.map_model_id("claude-sonnet-4-20250514")
assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
def test_us_region_maps_correctly(self):
"""US region should produce us.anthropic.* model IDs."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={
"claude-sonnet-4-20250514": "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",
},
):
backend = LiteLLMBackend(provider="bedrock", region="us-west-2")
result = backend.map_model_id("claude-sonnet-4-20250514")
assert result == "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0"
def test_fallback_for_unknown_model_in_eu(self):
"""Unknown models in EU should get eu.anthropic.* fallback, not bare 'bedrock/claude-...'."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="eu-west-1")
result = backend.map_model_id("claude-sonnet-4-20250514")
assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
def test_fallback_for_unknown_model_in_ap(self):
"""Unknown models in AP should get apac.anthropic.* fallback."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1")
result = backend.map_model_id("claude-3-5-haiku-20241022")
assert result == "bedrock/apac.anthropic.claude-3-5-haiku-20241022-v1:0"
def test_bedrock_format_passthrough(self):
"""Already-formatted Bedrock IDs should pass through unchanged."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="eu-central-1")
model = "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
result = backend.map_model_id(model)
assert result == model
def test_anthropic_dot_format_normalized(self):
"""Raw Bedrock IDs like 'anthropic.claude-...-v1:0' should normalize and map."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={
"claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0",
},
):
backend = LiteLLMBackend(provider="bedrock", region="eu-central-1")
result = backend.map_model_id("anthropic.claude-sonnet-4-20250514-v1:0")
assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
def test_region_prefixed_format_normalized(self):
"""'eu.anthropic.claude-...-v1:0' should normalize and map."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={
"claude-sonnet-4-20250514": "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0",
},
):
backend = LiteLLMBackend(provider="bedrock", region="eu-central-1")
result = backend.map_model_id("eu.anthropic.claude-sonnet-4-20250514-v1:0")
assert result == "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
def test_arn_passthrough(self):
"""Application inference profile ARNs must use the converse route."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-2")
arn = "arn:aws:bedrock:ap-southeast-2:123456789012:application-inference-profile/abc123"
result = backend.map_model_id(arn)
assert result == f"bedrock/converse/{arn}"
def test_ap_southeast_2_uses_au_prefix(self):
"""ap-southeast-2 (Sydney/Australia) should use 'au.' prefix, not 'apac.'."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-2")
result = backend.map_model_id("claude-sonnet-4-5-20250929")
assert result == "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0"
def test_override_pins_plain_name_to_app_profile_arn(self, monkeypatch):
"""HEADROOM_BEDROCK_MODEL_MAP pins a plain name to a specific app
profile ARN and routes it via the converse endpoint, winning over
the discovered map."""
arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66"
monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}")
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
# Discovery also has a system-defined sonnet-5; the override must win.
return_value={"claude-sonnet-5": "bedrock/global.anthropic.claude-sonnet-5"},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1")
assert backend.map_model_id("claude-sonnet-5") == f"bedrock/converse/{arn}"
def test_override_absent_falls_through_to_discovery(self, monkeypatch):
"""A model not pinned in the override map resolves via discovery."""
arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66"
monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}")
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={"claude-opus-4-8": "bedrock/global.anthropic.claude-opus-4-8"},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1")
assert backend.map_model_id("claude-opus-4-8") == (
"bedrock/global.anthropic.claude-opus-4-8"
)
def test_cross_region_prefixed_id_passes_through_directly(self):
"""au./us./eu./apac./global. prefixed IDs must be passed straight to
Bedrock without discovery remapping. Remapping can route to an
APPLICATION inference profile owned by another team, causing a 403.
Regression test for the bug reported in BUG_FIX.md."""
# Simulate a discovery map that contains an APPLICATION-type profile
# for the same underlying model (the bad case: wrong profile selected).
bad_app_profile = "bedrock/arn:aws:bedrock:ap-southeast-2:002037730852:application-inference-profile/6lgt8epqa0wf"
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={"claude-opus-4-8": bad_app_profile},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-2")
# The au. prefix must bypass discovery and pass through directly.
assert backend.map_model_id("au.anthropic.claude-opus-4-8") == (
"bedrock/au.anthropic.claude-opus-4-8"
)
def test_cross_region_prefixes_all_pass_through(self):
"""All five cross-region prefix families pass through without remapping."""
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={},
):
backend = LiteLLMBackend(provider="bedrock", region="us-east-1")
cases = [
("au.anthropic.claude-opus-4-8", "bedrock/au.anthropic.claude-opus-4-8"),
(
"us.anthropic.claude-sonnet-4-20250514-v1:0",
"bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",
),
(
"eu.anthropic.claude-3-5-sonnet-20241022-v2:0",
"bedrock/eu.anthropic.claude-3-5-sonnet-20241022-v2:0",
),
(
"apac.anthropic.claude-3-5-haiku-20241022-v1:0",
"bedrock/apac.anthropic.claude-3-5-haiku-20241022-v1:0",
),
("global.anthropic.claude-opus-4-8", "bedrock/global.anthropic.claude-opus-4-8"),
]
for model_in, expected in cases:
assert backend.map_model_id(model_in) == expected, (
f"Expected {expected!r} for input {model_in!r}"
)
def test_litellm_qualified_cross_region_id_passes_through_with_contaminated_map(self):
"""'bedrock/<cross-region-prefix>...' — the already LiteLLM-qualified form
documented in map_model_id's docstring — must also bypass discovery
remapping. Regression for the gap where only the bare 'us.anthropic...'
form was checked, so 'bedrock/us.anthropic...' still fell through to
normalization and could be remapped to a contaminating APPLICATION
profile in the discovery map."""
bad_app_profile = "bedrock/arn:aws:bedrock:ap-southeast-2:002037730852:application-inference-profile/6lgt8epqa0wf"
with patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value={"claude-opus-4-8": bad_app_profile},
):
backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-2")
assert backend.map_model_id("bedrock/au.anthropic.claude-opus-4-8") == (
"bedrock/au.anthropic.claude-opus-4-8"
)
# =============================================================================
# Normalize Bedrock Profile ID (edge cases)
# =============================================================================
class TestNormalizeBedrockProfileId:
"""Test normalization of various Bedrock profile ID formats."""
def test_eu_prefixed(self):
assert _normalize_bedrock_profile_id("eu.anthropic.claude-sonnet-4-20250514-v1:0") == (
"claude-sonnet-4-20250514"
)
def test_apac_prefixed(self):
assert _normalize_bedrock_profile_id("apac.anthropic.claude-3-5-sonnet-20241022-v2:0") == (
"claude-3-5-sonnet-20241022"
)
def test_us_prefixed(self):
assert _normalize_bedrock_profile_id("us.anthropic.claude-opus-4-20250514-v1:0") == (
"claude-opus-4-20250514"
)
def test_no_region_prefix(self):
assert _normalize_bedrock_profile_id("anthropic.claude-3-haiku-20240307-v1:0") == (
"claude-3-haiku-20240307"
)
def test_with_bedrock_slash_prefix(self):
assert (
_normalize_bedrock_profile_id("bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0")
== "claude-sonnet-4-20250514"
)
def test_non_claude_returns_none(self):
assert _normalize_bedrock_profile_id("eu.meta.llama-3-70b-v1:0") is None
def test_already_normalized(self):
assert _normalize_bedrock_profile_id("claude-sonnet-4-20250514") == (
"claude-sonnet-4-20250514"
)
def test_global_prefix_no_version_suffix(self):
# Current-gen cross-region profile: "global." prefix, no version
# suffix at all. Must strip the prefix and keep the bare name.
assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-8") == (
"claude-opus-4-8"
)
def test_global_prefix_bare_v_suffix(self):
# "global." prefix with an undated "-v1" (no ":revision") suffix.
assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-6-v1") == (
"claude-opus-4-6"
)
def test_global_prefix_dated_full_suffix(self):
# "global." prefix with the legacy dated "-vN:M" suffix.
assert (
_normalize_bedrock_profile_id("global.anthropic.claude-haiku-4-5-20251001-v1:0")
== "claude-haiku-4-5-20251001"
)
def test_global_prefix_next_gen_names(self):
assert (
_normalize_bedrock_profile_id("global.anthropic.claude-sonnet-5") == "claude-sonnet-5"
)
assert _normalize_bedrock_profile_id("global.anthropic.claude-fable-5") == "claude-fable-5"
# =============================================================================
# HEADROOM_BEDROCK_MODEL_MAP operator override parsing
# =============================================================================
class TestParseBedrockModelOverrides:
"""Test the HEADROOM_BEDROCK_MODEL_MAP override parser."""
def test_none_and_empty_yield_empty(self):
from headroom.backends.litellm import _parse_bedrock_model_overrides
assert _parse_bedrock_model_overrides(None) == {}
assert _parse_bedrock_model_overrides("") == {}
assert _parse_bedrock_model_overrides(" ") == {}
def test_single_pair(self):
from headroom.backends.litellm import _parse_bedrock_model_overrides
arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66"
assert _parse_bedrock_model_overrides(f"claude-sonnet-5={arn}") == {"claude-sonnet-5": arn}
def test_multiple_pairs_and_whitespace(self):
from headroom.backends.litellm import _parse_bedrock_model_overrides
raw = " claude-sonnet-5=arn:a , claude-opus-4-8=arn:b "
assert _parse_bedrock_model_overrides(raw) == {
"claude-sonnet-5": "arn:a",
"claude-opus-4-8": "arn:b",
}
def test_skips_malformed_entries(self):
from headroom.backends.litellm import _parse_bedrock_model_overrides
# Missing "=" and blank segments are skipped, valid pairs survive.
assert _parse_bedrock_model_overrides("garbage,,claude-sonnet-5=arn:a,=noname") == {
"claude-sonnet-5": "arn:a",
}
# =============================================================================
# Named profile forwarded to acompletion kwargs
# =============================================================================
_MODEL_MAP_US = {"claude-sonnet-4-20250514": "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0"}
_BODY = {
"model": "claude-sonnet-4-20250514",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 10,
}
def _make_fake_completion_resp():
mock_resp = MagicMock()
mock_resp.choices = [MagicMock()]
mock_resp.choices[0].message.content = "hello"
mock_resp.choices[0].message.tool_calls = None
mock_resp.choices[0].finish_reason = "stop"
mock_resp.usage.prompt_tokens = 10
mock_resp.usage.completion_tokens = 5
return mock_resp
class TestBedrockProfileForwardedToCompletion:
"""Regression: --bedrock-profile must be passed to acompletion(), not just to
_fetch_bedrock_inference_profiles() at startup. Without self.profile_name the
actual Bedrock call still uses ambient/default credentials even when the user
explicitly supplied a named SSO profile."""
def setup_method(self):
_bedrock_profiles_cache.clear()
async def test_send_message_passes_aws_profile_name(self):
"""send_message() must include aws_profile_name in the acompletion() kwargs."""
captured_kwargs: dict = {}
async def fake_acompletion(**kwargs):
captured_kwargs.update(kwargs)
return _make_fake_completion_resp()
with (
patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value=_MODEL_MAP_US,
),
patch("headroom.backends.litellm.acompletion", side_effect=fake_acompletion),
):
backend = LiteLLMBackend(
provider="bedrock", region="us-east-1", profile_name="my-sso-profile"
)
await backend.send_message(body=_BODY, headers={})
assert captured_kwargs.get("aws_profile_name") == "my-sso-profile"
async def test_stream_message_passes_aws_profile_name(self):
"""stream_message() must include aws_profile_name in the acompletion() kwargs."""
captured_kwargs: dict = {}
async def fake_acompletion(**kwargs):
captured_kwargs.update(kwargs)
async def _empty():
return
yield # pragma: no cover — makes this an async generator
return _empty()
with (
patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value=_MODEL_MAP_US,
),
patch("headroom.backends.litellm.acompletion", side_effect=fake_acompletion),
):
backend = LiteLLMBackend(
provider="bedrock", region="us-east-1", profile_name="my-sso-profile"
)
async for _ in backend.stream_message(body=_BODY, headers={}):
pass
assert captured_kwargs.get("aws_profile_name") == "my-sso-profile"
async def test_no_profile_does_not_set_aws_profile_name(self):
"""When no profile is configured, aws_profile_name must not appear in kwargs
(LiteLLM falls back to ambient credentials correctly)."""
captured_kwargs: dict = {}
async def fake_acompletion(**kwargs):
captured_kwargs.update(kwargs)
return _make_fake_completion_resp()
with (
patch(
"headroom.backends.litellm._fetch_bedrock_inference_profiles",
return_value=_MODEL_MAP_US,
),
patch("headroom.backends.litellm.acompletion", side_effect=fake_acompletion),
):
backend = LiteLLMBackend(provider="bedrock", region="us-east-1")
await backend.send_message(body=_BODY, headers={})
assert "aws_profile_name" not in captured_kwargs
Reference in a new issue View git blame Copy permalink