Skip to content

Commit 53df7e4

Browse files
mliu-clouderamliu-cloudera
committed
Mock Bedrock calls
1 parent ad85f43 commit 53df7e4

File tree

3 files changed

+201
-1
lines changed

3 files changed

+201
-1
lines changed

‎llm-service/app/tests/provider_mocks/__init__.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
#
2+
# CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
3+
# (C) Cloudera, Inc. 2025
4+
# All rights reserved.
5+
#
6+
# Applicable Open Source License: Apache 2.0
7+
#
8+
# NOTE: Cloudera open source products are modular software products
9+
# made up of hundreds of individual components, each of which was
10+
# individually copyrighted. Each Cloudera open source product is a
11+
# collective work under U.S. Copyright Law. Your license to use the
12+
# collective work is as provided in your written agreement with
13+
# Cloudera. Used apart from the collective work, this file is
14+
# licensed for your use pursuant to the open source license
15+
# identified above.
16+
#
17+
# This code is provided to you pursuant a written agreement with
18+
# (i) Cloudera, Inc. or (ii) a third-party authorized to distribute
19+
# this code. If you do not have a written agreement with Cloudera nor
20+
# with an authorized and properly licensed third party, you do not
21+
# have any rights to access nor to use this code.
22+
#
23+
# Absent a written agreement with Cloudera, Inc. ("Cloudera") to the
24+
# contrary, A) CLOUDERA PROVIDES THIS CODE TO YOU WITHOUT WARRANTIES OF ANY
25+
# KIND; (B) CLOUDERA DISCLAIMS ANY AND ALL EXPRESS AND IMPLIED
26+
# WARRANTIES WITH RESPECT TO THIS CODE, INCLUDING BUT NOT LIMITED TO
27+
# IMPLIED WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY AND
28+
# FITNESS FOR A PARTICULAR PURPOSE; (C) CLOUDERA IS NOT LIABLE TO YOU,
29+
# AND WILL NOT DEFEND, INDEMNIFY, NOR HOLD YOU HARMLESS FOR ANY CLAIMS
30+
# ARISING FROM OR RELATED TO THE CODE; AND (D)WITH RESPECT TO YOUR EXERCISE
31+
# OF ANY RIGHTS GRANTED TO YOU FOR THE CODE, CLOUDERA IS NOT LIABLE FOR ANY
32+
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, PUNITIVE OR
33+
# CONSEQUENTIAL DAMAGES INCLUDING, BUT NOT LIMITED TO, DAMAGES
34+
# RELATED TO LOST REVENUE, LOST PROFITS, LOSS OF INCOME, LOSS OF
35+
# BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
36+
# DATA.
37+
#
38+

‎llm-service/app/tests/provider_mocks/bedrock.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
#
2+
# CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
3+
# (C) Cloudera, Inc. 2025
4+
# All rights reserved.
5+
#
6+
# Applicable Open Source License: Apache 2.0
7+
#
8+
# NOTE: Cloudera open source products are modular software products
9+
# made up of hundreds of individual components, each of which was
10+
# individually copyrighted. Each Cloudera open source product is a
11+
# collective work under U.S. Copyright Law. Your license to use the
12+
# collective work is as provided in your written agreement with
13+
# Cloudera. Used apart from the collective work, this file is
14+
# licensed for your use pursuant to the open source license
15+
# identified above.
16+
#
17+
# This code is provided to you pursuant a written agreement with
18+
# (i) Cloudera, Inc. or (ii) a third-party authorized to distribute
19+
# this code. If you do not have a written agreement with Cloudera nor
20+
# with an authorized and properly licensed third party, you do not
21+
# have any rights to access nor to use this code.
22+
#
23+
# Absent a written agreement with Cloudera, Inc. ("Cloudera") to the
24+
# contrary, A) CLOUDERA PROVIDES THIS CODE TO YOU WITHOUT WARRANTIES OF ANY
25+
# KIND; (B) CLOUDERA DISCLAIMS ANY AND ALL EXPRESS AND IMPLIED
26+
# WARRANTIES WITH RESPECT TO THIS CODE, INCLUDING BUT NOT LIMITED TO
27+
# IMPLIED WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY AND
28+
# FITNESS FOR A PARTICULAR PURPOSE; (C) CLOUDERA IS NOT LIABLE TO YOU,
29+
# AND WILL NOT DEFEND, INDEMNIFY, NOR HOLD YOU HARMLESS FOR ANY CLAIMS
30+
# ARISING FROM OR RELATED TO THE CODE; AND (D)WITH RESPECT TO YOUR EXERCISE
31+
# OF ANY RIGHTS GRANTED TO YOU FOR THE CODE, CLOUDERA IS NOT LIABLE FOR ANY
32+
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, PUNITIVE OR
33+
# CONSEQUENTIAL DAMAGES INCLUDING, BUT NOT LIMITED TO, DAMAGES
34+
# RELATED TO LOST REVENUE, LOST PROFITS, LOSS OF INCOME, LOSS OF
35+
# BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
36+
# DATA.
37+
#
38+
from typing import Generator
39+
from unittest.mock import patch
40+
from urllib.parse import urljoin
41+
42+
import botocore
43+
import pytest
44+
import responses
45+
46+
from app.config import settings
47+
48+
49+
@pytest.fixture
50+
def mock_bedrock() -> Generator[None, None, None]:
51+
BEDROCK_URL_BASE = f"https://bedrock.{settings.aws_default_region}.amazonaws.com/"
52+
TEXT_MODELS = [
53+
("test.unavailable-text-model-v1", "NOT_AVAILABLE"),
54+
("test.available-text-model-v1", "AVAILABLE"),
55+
]
56+
EMBEDDING_MODELS = [
57+
("test.unavailable-embedding-model-v1", "NOT_AVAILABLE"),
58+
("test.available-embedding-model-v1", "AVAILABLE"),
59+
]
60+
61+
r_mock = responses.RequestsMock(assert_all_requests_are_fired=False)
62+
for model_id, availability in TEXT_MODELS + EMBEDDING_MODELS:
63+
r_mock.get(
64+
urljoin(
65+
BEDROCK_URL_BASE,
66+
f"foundation-model-availability/{model_id}:0",
67+
),
68+
json={
69+
"agreementAvailability": {
70+
"errorMessage": None,
71+
"status": availability,
72+
},
73+
"authorizationStatus": "AUTHORIZED",
74+
"entitlementAvailability": availability,
75+
"modelId": model_id,
76+
"regionAvailability": "AVAILABLE",
77+
},
78+
)
79+
80+
make_api_call = botocore.client.BaseClient._make_api_call
81+
82+
def mock_make_api_call(self, operation_name: str, api_params: dict[str, str]):
83+
"""Mock Bedrock calls, since moto doesn't have full coverage.
84+
85+
Based on https://docs.getmoto.org/en/latest/docs/services/patching_other_services.html.
86+
87+
"""
88+
if operation_name == "ListFoundationModels":
89+
modality = api_params["byOutputModality"]
90+
if modality == "TEXT":
91+
return {
92+
"modelSummaries": [
93+
{
94+
"modelArn": f"arn:aws:bedrock:{settings.aws_default_region}::foundation-model/{model_id}:0",
95+
"modelId": f"{model_id}:0",
96+
"modelName": model_id.upper(),
97+
"providerName": "Test",
98+
"inputModalities": ["TEXT"],
99+
"outputModalities": ["TEXT"],
100+
"responseStreamingSupported": True,
101+
"customizationsSupported": [],
102+
"inferenceTypesSupported": ["ON_DEMAND"],
103+
"modelLifecycle": {"status": "ACTIVE"},
104+
}
105+
for model_id, _ in TEXT_MODELS
106+
],
107+
}
108+
elif modality == "EMBEDDING":
109+
return {
110+
"modelSummaries": [
111+
{
112+
"modelArn": f"arn:aws:bedrock:{settings.aws_default_region}::foundation-model/{model_id}:0",
113+
"modelId": f"{model_id}:0",
114+
"modelName": model_id.upper(),
115+
"providerName": "Test",
116+
"inputModalities": ["TEXT"],
117+
"outputModalities": ["EMBEDDING"],
118+
"responseStreamingSupported": False,
119+
"customizationsSupported": [],
120+
"inferenceTypesSupported": ["ON_DEMAND"],
121+
"modelLifecycle": {"status": "ACTIVE"},
122+
}
123+
for model_id, _ in EMBEDDING_MODELS
124+
],
125+
}
126+
else:
127+
raise ValueError(f"test encountered unexpected modality {modality}")
128+
elif operation_name == "ListInferenceProfiles":
129+
return {
130+
"inferenceProfileSummaries": [
131+
{
132+
"inferenceProfileName": f"US {model_id.upper()}",
133+
"description": f"Routes requests to {model_id.upper()} in {settings.aws_default_region}.",
134+
"inferenceProfileArn": f"arn:aws:bedrock:{settings.aws_default_region}:123456789012:inference-profile/{model_id}:0",
135+
"models": [
136+
{
137+
"modelArn": f"arn:aws:bedrock:{settings.aws_default_region}::foundation-model/{model_id}:0"
138+
},
139+
],
140+
"inferenceProfileId": f"{model_id}:0",
141+
"status": "ACTIVE",
142+
"type": "SYSTEM_DEFINED",
143+
}
144+
for model_id, _ in TEXT_MODELS + EMBEDDING_MODELS
145+
],
146+
}
147+
148+
else:
149+
# passthrough
150+
return make_api_call(self, operation_name, api_params)
151+
152+
with patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call):
153+
with r_mock:
154+
yield
155+
156+
157+
def test_bedrock(mock_bedrock) -> None:
158+
from app.services.models.providers import BedrockModelProvider
159+
160+
BedrockModelProvider.list_available_models()
161+
BedrockModelProvider._get_model_arns()

‎llm-service/pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,8 @@ dev = [
7070
"hypothesis>=6.122.1",
7171
"black>=24.10.0",
7272
"types-boto3-s3>=1.36.1",
73-
"types-requests>=2.32.0"
73+
"types-requests>=2.32.0",
74+
"responses>=0.25.8",
7475
]
7576

7677
[tool.pytest.ini_options]

0 commit comments

Comments
 (0)