@@ -129,7 +129,7 @@ def test__init__with_default_region(session_cls, mock_client_method):
129129 with unittest .mock .patch .object (os , "environ" , {}):
130130 BedrockModel ()
131131 session_cls .return_value .client .assert_called_with (
132- region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY
132+ region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY , endpoint_url = None
133133 )
134134
135135
@@ -139,22 +139,22 @@ def test__init__with_session_region(session_cls, mock_client_method):
139139
140140 BedrockModel ()
141141
142- mock_client_method .assert_called_with (region_name = "eu-blah-1" , config = ANY , service_name = ANY )
142+ mock_client_method .assert_called_with (region_name = "eu-blah-1" , config = ANY , service_name = ANY , endpoint_url = None )
143143
144144
145145def test__init__with_custom_region (mock_client_method ):
146146 """Test that BedrockModel uses the provided region."""
147147 custom_region = "us-east-1"
148148 BedrockModel (region_name = custom_region )
149- mock_client_method .assert_called_with (region_name = custom_region , config = ANY , service_name = ANY )
149+ mock_client_method .assert_called_with (region_name = custom_region , config = ANY , service_name = ANY , endpoint_url = None )
150150
151151
152152def test__init__with_default_environment_variable_region (mock_client_method ):
153153 """Test that BedrockModel uses the AWS_REGION since we code that in."""
154154 with unittest .mock .patch .object (os , "environ" , {"AWS_REGION" : "eu-west-2" }):
155155 BedrockModel ()
156156
157- mock_client_method .assert_called_with (region_name = "eu-west-2" , config = ANY , service_name = ANY )
157+ mock_client_method .assert_called_with (region_name = "eu-west-2" , config = ANY , service_name = ANY , endpoint_url = None )
158158
159159
160160def test__init__region_precedence (mock_client_method , session_cls ):
@@ -164,21 +164,38 @@ def test__init__region_precedence(mock_client_method, session_cls):
164164
165165 # specifying a region always wins out
166166 BedrockModel (region_name = "us-specified-1" )
167- mock_client_method .assert_called_with (region_name = "us-specified-1" , config = ANY , service_name = ANY )
167+ mock_client_method .assert_called_with (
168+ region_name = "us-specified-1" , config = ANY , service_name = ANY , endpoint_url = None
169+ )
168170
169171 # other-wise uses the session's
170172 BedrockModel ()
171- mock_client_method .assert_called_with (region_name = "us-session-1" , config = ANY , service_name = ANY )
173+ mock_client_method .assert_called_with (
174+ region_name = "us-session-1" , config = ANY , service_name = ANY , endpoint_url = None
175+ )
172176
173177 # environment variable next
174178 session_cls .return_value .region_name = None
175179 BedrockModel ()
176- mock_client_method .assert_called_with (region_name = "us-environment-1" , config = ANY , service_name = ANY )
180+ mock_client_method .assert_called_with (
181+ region_name = "us-environment-1" , config = ANY , service_name = ANY , endpoint_url = None
182+ )
177183
178184 mock_os_environ .pop ("AWS_REGION" )
179185 session_cls .return_value .region_name = None # No session region
180186 BedrockModel ()
181- mock_client_method .assert_called_with (region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY )
187+ mock_client_method .assert_called_with (
188+ region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY , endpoint_url = None
189+ )
190+
191+
192+ def test__init__with_endpoint_url (mock_client_method ):
193+ """Test that BedrockModel uses the provided endpoint_url for VPC endpoints."""
194+ custom_endpoint = "https://vpce-12345-abcde.bedrock-runtime.us-west-2.vpce.amazonaws.com"
195+ BedrockModel (endpoint_url = custom_endpoint )
196+ mock_client_method .assert_called_with (
197+ region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY , endpoint_url = custom_endpoint
198+ )
182199
183200
184201def test__init__with_region_and_session_raises_value_error ():
0 commit comments