From eed154f21d50571f216f9a389737756310d96b2b Mon Sep 17 00:00:00 2001 From: Sean Patch Date: Tue, 3 Feb 2026 11:31:03 -0500 Subject: [PATCH] fix(identity): update endpoint for Create/UpdateWorkloadIdentity --- src/bedrock_agentcore/services/identity.py | 7 +--- .../services/test_identity.py | 41 ++++++++----------- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/src/bedrock_agentcore/services/identity.py b/src/bedrock_agentcore/services/identity.py index 80d30ac..5ffc869 100644 --- a/src/bedrock_agentcore/services/identity.py +++ b/src/bedrock_agentcore/services/identity.py @@ -78,9 +78,6 @@ def __init__(self, region: str): self.cp_client = boto3.client( "bedrock-agentcore-control", region_name=region, endpoint_url=get_control_plane_endpoint(region) ) - self.identity_client = boto3.client( - "bedrock-agentcore-control", region_name=region, endpoint_url=get_data_plane_endpoint(region) - ) self.dp_client = boto3.client( "bedrock-agentcore", region_name=region, endpoint_url=get_data_plane_endpoint(region) ) @@ -122,7 +119,7 @@ def create_workload_identity( self.logger.info("Creating workload identity...") if not name: name = f"workload-{uuid.uuid4().hex[:8]}" - return self.identity_client.create_workload_identity( + return self.cp_client.create_workload_identity( name=name, allowedResourceOauth2ReturnUrls=allowed_resource_oauth_2_return_urls or [] ) @@ -131,7 +128,7 @@ def update_workload_identity(self, name: str, allowed_resource_oauth_2_return_ur self.logger.info( "Updating workload identity '%s' with callback urls: %s", name, allowed_resource_oauth_2_return_urls ) - return self.identity_client.update_workload_identity( + return self.cp_client.update_workload_identity( name=name, allowedResourceOauth2ReturnUrls=allowed_resource_oauth_2_return_urls ) diff --git a/tests/bedrock_agentcore/services/test_identity.py b/tests/bedrock_agentcore/services/test_identity.py index d3801c1..918c61d 100644 --- a/tests/bedrock_agentcore/services/test_identity.py +++ b/tests/bedrock_agentcore/services/test_identity.py @@ -376,9 +376,8 @@ def test_get_workload_access_token_with_user_token(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -406,9 +405,8 @@ def test_get_workload_access_token_with_user_id(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -435,9 +433,8 @@ def test_get_workload_access_token_without_user_info(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -461,28 +458,27 @@ def test_create_workload_identity(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) # Test with provided name custom_name = "my-custom-workload" expected_response = {"name": custom_name, "workloadIdentityId": "workload-123"} - mock_identity_client.create_workload_identity.return_value = expected_response + mock_cp_client.create_workload_identity.return_value = expected_response result = identity_client.create_workload_identity(name=custom_name) assert result == expected_response - mock_identity_client.create_workload_identity.assert_called_with( + mock_cp_client.create_workload_identity.assert_called_with( name=custom_name, allowedResourceOauth2ReturnUrls=[] ) # Test without provided name (auto-generated) - mock_identity_client.reset_mock() + mock_cp_client.reset_mock() expected_response_auto = {"name": "workload-abcd1234", "workloadIdentityId": "workload-456"} - mock_identity_client.create_workload_identity.return_value = expected_response_auto + mock_cp_client.create_workload_identity.return_value = expected_response_auto with patch("uuid.uuid4") as mock_uuid: mock_uuid.return_value.hex = "abcd1234efgh5678" @@ -490,7 +486,7 @@ def test_create_workload_identity(self): result = identity_client.create_workload_identity() assert result == expected_response_auto - mock_identity_client.create_workload_identity.assert_called_with( + mock_cp_client.create_workload_identity.assert_called_with( name="workload-abcd1234", allowedResourceOauth2ReturnUrls=[] ) @@ -499,9 +495,8 @@ def test_update_workload_identity(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -509,12 +504,12 @@ def test_update_workload_identity(self): allowed_urls = ["https://unit-test.com/callback", "https://test.com/oauth"] expected_response = {"name": workload_name, "allowedResourceOauth2ReturnUrls": allowed_urls} - mock_identity_client.update_workload_identity.return_value = expected_response + mock_cp_client.update_workload_identity.return_value = expected_response result = identity_client.update_workload_identity(workload_name, allowed_urls) assert result == expected_response - mock_identity_client.update_workload_identity.assert_called_once_with( + mock_cp_client.update_workload_identity.assert_called_once_with( name=workload_name, allowedResourceOauth2ReturnUrls=allowed_urls ) @@ -523,9 +518,8 @@ def test_get_workload_identity(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -545,9 +539,8 @@ def test_complete_resource_token_auth_with_user_id(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -570,9 +563,8 @@ def test_complete_resource_token_auth_with_user_token(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region) @@ -595,9 +587,8 @@ def test_complete_resource_token_auth_with_invalid_identifier(self): with patch("boto3.client") as mock_boto_client: mock_cp_client = Mock() - mock_identity_client = Mock() mock_dp_client = Mock() - mock_boto_client.side_effect = [mock_cp_client, mock_identity_client, mock_dp_client] + mock_boto_client.side_effect = [mock_cp_client, mock_dp_client] identity_client = IdentityClient(region)