scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

test(bedrock): wrap long lines flagged by review in migrated unit tests

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng 2026-09-20 09:37:17 +00:00
parent feca00248c
commit ba629f2537
5 changed files with 34 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py
View file

@ -11,7 +11,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None:
with (
patch(
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params"
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration."
"_get_boto_credentials_from_optional_params"
),
patch(
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers"
@ -31,7 +32,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None:
assert (
request.endpoint_url
== "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012%3Aapplication-inference-profile%2Fabcdefghi123/invoke"
== "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012"
"%3Aapplication-inference-profile%2Fabcdefghi123/invoke"
)
@ -41,7 +43,8 @@ def test_bedrock_image_prepare_request_without_arn() -> None:
with (
patch(
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params"
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration."
"_get_boto_credentials_from_optional_params"
),
patch(
"litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers"

13
tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py
View file

@ -419,7 +419,9 @@ def test_bedrock_passthrough_model_id_arn_encoding():
), f"ARN slash should be encoded, but found unencoded version in: {url_str}"
# Verify the complete expected URL structure
expected_encoded_model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7"
expected_encoded_model_id = (
"arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7"
)
expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/converse"
assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}"
@ -515,7 +517,10 @@ def test_bedrock_passthrough_model_id_without_arn():
def _event_frame(event_type: str, payload: dict) -> bytes:
def header(name: str, value: str) -> bytes:
name_b, value_b = name.encode(), value.encode()
return struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b
return (
struct.pack("!B", len(name_b)) + name_b
+ struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b
)
payload_b = json.dumps(payload, separators=(",", ":")).encode()
headers_b = (
@ -589,7 +594,9 @@ def _feed(collector: PassthroughStreamCollector, stream: bytes, chunk_size: int
def test_converse_stream_collector_keeps_usage_without_retaining_the_stream():
texts = [f"tok{i} " for i in range(4000)]
stream = _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000)
stream = (
_event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000)
)
_feed(_converse_stream_collector(), stream)
tracemalloc.start()

8
tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py
View file

@ -23,7 +23,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path):
monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials"))
monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config"))
monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"):
for env_var in (
"AWS_PROFILE",
"AWS_DEFAULT_PROFILE",
"AWS_BEARER_TOKEN_BEDROCK",
"AWS_REGION_NAME",
"AWS_DEFAULT_REGION",
):
monkeypatch.delenv(env_var, raising=False)

11
tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
View file

@ -21,7 +21,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path):
monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials"))
monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config"))
monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"):
for env_var in (
"AWS_PROFILE",
"AWS_DEFAULT_PROFILE",
"AWS_BEARER_TOKEN_BEDROCK",
"AWS_REGION_NAME",
"AWS_DEFAULT_REGION",
):
monkeypatch.delenv(env_var, raising=False)
# Mock response for Bedrock rerank
@ -39,7 +45,8 @@ bedrock_rerank_response = {
test_query = "What is the capital of the United States?"
test_documents = [
"Carson City is the capital city of the American state of Nevada.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. "
"Its capital is Saipan.",
"Washington, D.C. is the capital of the United States.",
]

3
tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py
View file

@ -46,7 +46,8 @@ def test_transform_search_request_encodes_vector_store_id():
assert (
url
== "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve"
== "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother"
"%3Fx%3D1%23frag/retrieve"
)
assert body["retrievalQuery"].get("text") == "hello"