@@ -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,28 @@ 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 (region_name = "us-specified-1" , config = ANY , service_name = ANY , endpoint_url = None )
168168
169169 # other-wise uses the session's
170170 BedrockModel ()
171- mock_client_method .assert_called_with (region_name = "us-session-1" , config = ANY , service_name = ANY )
171+ mock_client_method .assert_called_with (region_name = "us-session-1" , config = ANY , service_name = ANY , endpoint_url = None )
172172
173173 # environment variable next
174174 session_cls .return_value .region_name = None
175175 BedrockModel ()
176- mock_client_method .assert_called_with (region_name = "us-environment-1" , config = ANY , service_name = ANY )
176+ mock_client_method .assert_called_with (region_name = "us-environment-1" , config = ANY , service_name = ANY , endpoint_url = None )
177177
178178 mock_os_environ .pop ("AWS_REGION" )
179179 session_cls .return_value .region_name = None # No session region
180180 BedrockModel ()
181- mock_client_method .assert_called_with (region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY )
181+ mock_client_method .assert_called_with (region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY , endpoint_url = None )
182+
183+
184+ def test__init__with_endpoint_url (mock_client_method ):
185+ """Test that BedrockModel uses the provided endpoint_url for VPC endpoints."""
186+ custom_endpoint = "https://vpce-12345-abcde.bedrock-runtime.us-west-2.vpce.amazonaws.com"
187+ BedrockModel (endpoint_url = custom_endpoint )
188+ mock_client_method .assert_called_with (region_name = DEFAULT_BEDROCK_REGION , config = ANY , service_name = ANY , endpoint_url = custom_endpoint )
182189
183190
184191def test__init__with_region_and_session_raises_value_error ():
0 commit comments