adjusting cassettes for bedrock

This commit is contained in:
lorenzejay
2026-05-26 10:36:42 -07:00
parent d883e6e587
commit 4e93aa6660
2 changed files with 38 additions and 6 deletions

View File

@@ -5,6 +5,7 @@ from collections.abc import Generator
import gzip
import os
from pathlib import Path
import re
import tempfile
from typing import Any
@@ -20,8 +21,25 @@ except ModuleNotFoundError:
env_test_path = Path(__file__).parent / ".env.test"
load_dotenv(env_test_path, override=True)
load_dotenv(override=True)
load_dotenv(env_test_path, override=False)
load_dotenv(override=False)
BEDROCK_HOST_PLACEHOLDER = "bedrock-runtime.vcr.amazonaws.com"
_BEDROCK_HOST_RE = re.compile(r"^bedrock-runtime\.[a-z0-9-]+\.amazonaws\.com$")
def _normalize_bedrock_host(host: str) -> str:
if _BEDROCK_HOST_RE.match(host):
return BEDROCK_HOST_PLACEHOLDER
return host
def bedrock_host_matcher(r1: Request, r2: Request) -> bool: # type: ignore[no-any-unimported]
"""Match Bedrock requests across AWS regions (CI uses us-east-1, local may use us-west-2)."""
return _normalize_bedrock_host(r1.host or "") == _normalize_bedrock_host(
r2.host or ""
)
def _patched_make_vcr_request(httpx_request: Any, **kwargs: Any) -> Any:
@@ -213,6 +231,10 @@ def _filter_request_headers(request: Request) -> Request: # type: ignore[no-any
placeholder_host = "fake-azure-endpoint.openai.azure.com"
request.uri = request.uri.replace(original_host, placeholder_host)
# Normalize Bedrock regional endpoints so cassettes work in any AWS region.
if request.host and _BEDROCK_HOST_RE.match(request.host):
request.uri = request.uri.replace(request.host, BEDROCK_HOST_PLACEHOLDER)
return request
@@ -230,6 +252,11 @@ def _filter_response_headers(response: dict[str, Any]) -> dict[str, Any] | None:
if body == "" or body == b"" or content_length == ["0"]:
return None
status_code = response.get("status", {}).get("code")
if isinstance(status_code, int) and status_code >= 400:
# Avoid persisting auth/model errors when re-recording without valid AWS creds.
return None
for encoding_header in ["Content-Encoding", "content-encoding"]:
if encoding_header in headers:
encoding = headers.pop(encoding_header)
@@ -280,6 +307,11 @@ def vcr_cassette_dir(request: Any) -> str:
return str(cassette_dir)
def pytest_recording_configure(vcr: Any, config: Any) -> None:
"""Register custom VCR matchers for each test cassette session."""
vcr.register_matcher("bedrock_host", bedrock_host_matcher)
@pytest.fixture(scope="module")
def vcr_config(vcr_cassette_dir: str) -> dict[str, Any]:
"""Configure VCR with organized cassette storage."""
@@ -290,7 +322,7 @@ def vcr_config(vcr_cassette_dir: str) -> dict[str, Any]:
"before_record_request": _filter_request_headers,
"before_record_response": _filter_response_headers,
"filter_query_parameters": ["key"],
"match_on": ["method", "scheme", "host", "port", "path"],
"match_on": ["method", "scheme", "bedrock_host", "port", "path"],
}
if os.getenv("GITHUB_ACTIONS") == "true":

View File

@@ -24,10 +24,10 @@ interactions:
x-amz-security-token:
- X-AMZ-SECURITY-TOKEN-XXX
method: POST
uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-6/converse
uri: https://bedrock-runtime.vcr.amazonaws.com/model/us.anthropic.claude-sonnet-4-6/converse
response:
body:
string: '{"metrics":{"latencyMs":2109},"output":{"message":{"content":[{"text":"Blank"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"cacheReadInputTokenCount":0,"cacheReadInputTokens":0,"cacheWriteInputTokenCount":0,"cacheWriteInputTokens":0,"inputTokens":1606,"outputTokens":5,"serverToolUsage":{},"totalTokens":1611}}'
string: '{"metrics":{"latencyMs":2732},"output":{"message":{"content":[{"text":"Blank"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"cacheReadInputTokenCount":0,"cacheReadInputTokens":0,"cacheWriteInputTokenCount":0,"cacheWriteInputTokens":0,"inputTokens":1606,"outputTokens":5,"serverToolUsage":{},"totalTokens":1611}}'
headers:
Connection:
- keep-alive
@@ -36,7 +36,7 @@ interactions:
Content-Type:
- application/json
Date:
- Tue, 26 May 2026 17:06:30 GMT
- Tue, 26 May 2026 17:34:51 GMT
x-amzn-RequestId:
- X-AMZN-REQUESTID-XXX
status: