generated from amazon-archives/__template_Apache-2.0
-
Notifications
You must be signed in to change notification settings - Fork 100
feat: add ResourcePolicyClient for resource-based policy management #328
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
+272
−0
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| """Client for managing resource-based policies on Bedrock AgentCore resources.""" | ||
|
|
||
| import json | ||
| import logging | ||
| from typing import Optional, Union | ||
|
|
||
| import boto3 | ||
|
|
||
| from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint | ||
|
|
||
|
|
||
| class ResourcePolicyClient: | ||
| """Client for managing resource-based policies on Bedrock AgentCore resources. | ||
|
|
||
| Resource-based policies control which principals can invoke and manage | ||
| Agent Runtime, Endpoint, and Gateway resources. | ||
| """ | ||
|
|
||
| def __init__(self, region: str): | ||
| """Initialize the client for the specified region.""" | ||
| self.region = region | ||
| self.client = boto3.client( | ||
| "bedrock-agentcore-control", | ||
| region_name=region, | ||
| endpoint_url=get_control_plane_endpoint(region), | ||
| ) | ||
| self.logger = logging.getLogger("bedrock_agentcore.resource_policy_client") | ||
|
|
||
| def put_resource_policy(self, resource_arn: str, policy: Union[str, dict]) -> dict: | ||
| """Create or update a resource-based policy. | ||
|
|
||
| Args: | ||
| resource_arn: ARN of the resource to attach the policy to. | ||
| policy: Policy document as a dict (auto-serialized) or JSON string. | ||
|
|
||
| Returns: | ||
| The stored policy as a dict. | ||
| """ | ||
| policy_str = json.dumps(policy) if isinstance(policy, dict) else policy | ||
| self.logger.info("Putting resource policy for %s", resource_arn) | ||
| resp = self.client.put_resource_policy(resourceArn=resource_arn, policy=policy_str) | ||
| return json.loads(resp["policy"]) | ||
|
|
||
| def get_resource_policy(self, resource_arn: str) -> Optional[dict]: | ||
| """Get the resource-based policy for a resource. | ||
|
|
||
| Args: | ||
| resource_arn: ARN of the resource. | ||
|
|
||
| Returns: | ||
| The policy as a dict, or None if no policy is attached. | ||
| """ | ||
| self.logger.info("Getting resource policy for %s", resource_arn) | ||
| resp = self.client.get_resource_policy(resourceArn=resource_arn) | ||
| return json.loads(resp["policy"]) if "policy" in resp else None | ||
|
|
||
| def delete_resource_policy(self, resource_arn: str) -> dict: | ||
| """Delete the resource-based policy from a resource. | ||
|
|
||
| Args: | ||
| resource_arn: ARN of the resource. | ||
|
|
||
| Returns: | ||
| Raw boto3 response. | ||
|
|
||
| Raises: | ||
| ClientError: ResourceNotFoundException if no policy exists. | ||
| """ | ||
| self.logger.info("Deleting resource policy for %s", resource_arn) | ||
| return self.client.delete_resource_policy(resourceArn=resource_arn) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| """Unit tests for ResourcePolicyClient.""" | ||
|
|
||
| import json | ||
| from unittest.mock import Mock, patch | ||
|
|
||
| from bedrock_agentcore.services.resource_policy import ResourcePolicyClient | ||
|
|
||
| TEST_REGION = "us-east-1" | ||
| TEST_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-agent" | ||
|
|
||
|
|
||
| class TestResourcePolicyClient: | ||
| def test_initialization(self): | ||
| with patch("boto3.client") as mock_boto: | ||
| client = ResourcePolicyClient(region=TEST_REGION) | ||
| assert client.region == TEST_REGION | ||
| mock_boto.assert_called_once_with( | ||
| "bedrock-agentcore-control", | ||
| region_name=TEST_REGION, | ||
| endpoint_url=f"https://bedrock-agentcore-control.{TEST_REGION}.amazonaws.com", | ||
| ) | ||
|
|
||
| def test_put_serializes_dict_to_json(self): | ||
| with patch("boto3.client") as mock_boto: | ||
| mock_client = Mock() | ||
| mock_boto.return_value = mock_client | ||
| policy = {"Version": "2012-10-17", "Statement": []} | ||
| mock_client.put_resource_policy.return_value = {"policy": json.dumps(policy)} | ||
|
|
||
| client = ResourcePolicyClient(region=TEST_REGION) | ||
| result = client.put_resource_policy(TEST_ARN, policy) | ||
|
|
||
| mock_client.put_resource_policy.assert_called_once_with(resourceArn=TEST_ARN, policy=json.dumps(policy)) | ||
| assert result == policy | ||
|
|
||
| def test_get_returns_none_when_no_policy(self): | ||
| with patch("boto3.client") as mock_boto: | ||
| mock_client = Mock() | ||
| mock_boto.return_value = mock_client | ||
| mock_client.get_resource_policy.return_value = {} | ||
|
|
||
| client = ResourcePolicyClient(region=TEST_REGION) | ||
| assert client.get_resource_policy(TEST_ARN) is None | ||
|
|
||
| def test_get_deserializes_policy(self): | ||
| with patch("boto3.client") as mock_boto: | ||
| mock_client = Mock() | ||
| mock_boto.return_value = mock_client | ||
| policy = {"Version": "2012-10-17"} | ||
| mock_client.get_resource_policy.return_value = {"policy": json.dumps(policy)} | ||
|
|
||
| client = ResourcePolicyClient(region=TEST_REGION) | ||
| assert client.get_resource_policy(TEST_ARN) == policy | ||
|
|
||
| def test_delete_passes_through(self): | ||
| with patch("boto3.client") as mock_boto: | ||
| mock_client = Mock() | ||
| mock_boto.return_value = mock_client | ||
| mock_client.delete_resource_policy.return_value = {} | ||
|
|
||
| client = ResourcePolicyClient(region=TEST_REGION) | ||
| client.delete_resource_policy(TEST_ARN) | ||
| mock_client.delete_resource_policy.assert_called_once_with(resourceArn=TEST_ARN) |
Empty file.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,132 @@ | ||
| """Integration tests for ResourcePolicyClient. | ||
|
|
||
| E2e tests against real AWS resources. No mocking. | ||
|
|
||
| Requires: | ||
| RESOURCE_POLICY_TEST_ARN: ARN of an Agent Runtime to attach policies to | ||
| RESOURCE_POLICY_TEST_PRINCIPAL: IAM principal ARN to use in test policies | ||
| BEDROCK_TEST_REGION: AWS region (default: us-west-2) | ||
|
|
||
| Run: pytest -xvs tests_integ/services/test_resource_policy.py | ||
| """ | ||
|
|
||
| import json | ||
| import os | ||
|
|
||
| import pytest | ||
| from botocore.exceptions import ClientError | ||
|
|
||
| from bedrock_agentcore.services.resource_policy import ResourcePolicyClient | ||
|
|
||
|
|
||
| def _make_policy( | ||
| resource_arn: str, | ||
| principal_arn: str, | ||
| action: str = "bedrock-agentcore:InvokeAgentRuntime", | ||
| ) -> dict: | ||
| """Build a minimal valid policy dict for testing.""" | ||
| return { | ||
| "Version": "2012-10-17", | ||
| "Statement": [ | ||
| { | ||
| "Effect": "Allow", | ||
| "Principal": {"AWS": principal_arn}, | ||
| "Action": action, | ||
| "Resource": resource_arn, | ||
| } | ||
| ], | ||
| } | ||
|
|
||
|
|
||
| @pytest.mark.integration | ||
| class TestResourcePolicyClient: | ||
| @classmethod | ||
| def setup_class(cls): | ||
| cls.resource_arn = os.environ.get("RESOURCE_POLICY_TEST_ARN") | ||
| cls.principal_arn = os.environ.get("RESOURCE_POLICY_TEST_PRINCIPAL") | ||
| cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-west-2") | ||
|
|
||
| if not cls.resource_arn: | ||
| pytest.fail("RESOURCE_POLICY_TEST_ARN env var is required") | ||
| if not cls.principal_arn: | ||
| pytest.fail("RESOURCE_POLICY_TEST_PRINCIPAL env var is required") | ||
|
|
||
| cls.client = ResourcePolicyClient(region=cls.region) | ||
|
|
||
| def setup_method(self): | ||
| """Runs before each test — ensures no policy is attached.""" | ||
| try: | ||
| self.client.delete_resource_policy(self.resource_arn) | ||
| except Exception: | ||
| pass | ||
|
|
||
| @classmethod | ||
| def teardown_class(cls): | ||
| """Remove any policy left by the last test so we don't leave side effects on the resource.""" | ||
| try: | ||
| cls.client.delete_resource_policy(cls.resource_arn) | ||
| except Exception: | ||
| pass | ||
|
|
||
| @pytest.mark.order(1) | ||
| def test_get_returns_none_when_no_policy(self): | ||
| """get on a resource with no policy returns None.""" | ||
| result = self.client.get_resource_policy(self.resource_arn) | ||
| assert result is None | ||
|
|
||
| @pytest.mark.order(2) | ||
| def test_put_get_round_trip(self): | ||
| """put(policy) then get() returns matching policy as a dict.""" | ||
| policy = _make_policy(self.resource_arn, self.principal_arn) | ||
|
|
||
| put_result = self.client.put_resource_policy(self.resource_arn, policy) | ||
| assert isinstance(put_result, dict) | ||
| assert put_result["Version"] == policy["Version"] | ||
|
|
||
| result = self.client.get_resource_policy(self.resource_arn) | ||
| assert isinstance(result, dict) | ||
| assert result["Version"] == policy["Version"] | ||
| assert result["Statement"][0]["Effect"] == "Allow" | ||
| assert result["Statement"][0]["Resource"] == self.resource_arn | ||
|
|
||
| @pytest.mark.order(3) | ||
| def test_put_overwrites(self): | ||
| """put(A) then put(B) then get() returns B.""" | ||
| policy_a = _make_policy(self.resource_arn, self.principal_arn, action="bedrock-agentcore:InvokeAgentRuntime") | ||
| policy_b = _make_policy(self.resource_arn, self.principal_arn, action="bedrock-agentcore:GetAgentCard") | ||
|
|
||
| self.client.put_resource_policy(self.resource_arn, policy_a) | ||
| self.client.put_resource_policy(self.resource_arn, policy_b) | ||
| result = self.client.get_resource_policy(self.resource_arn) | ||
|
|
||
| assert result["Statement"][0]["Action"] == "bedrock-agentcore:GetAgentCard" | ||
|
|
||
| @pytest.mark.order(4) | ||
| def test_delete_removes_policy(self): | ||
| """put(policy) then delete() then get() returns None.""" | ||
| policy = _make_policy(self.resource_arn, self.principal_arn) | ||
| self.client.put_resource_policy(self.resource_arn, policy) | ||
| self.client.delete_resource_policy(self.resource_arn) | ||
|
|
||
| result = self.client.get_resource_policy(self.resource_arn) | ||
| assert result is None | ||
|
|
||
| @pytest.mark.order(5) | ||
| def test_delete_on_no_policy_raises(self): | ||
| """delete on a resource with no policy raises ResourceNotFoundException.""" | ||
| with pytest.raises(ClientError) as exc_info: | ||
| self.client.delete_resource_policy(self.resource_arn) | ||
| assert exc_info.value.response["Error"]["Code"] == "ResourceNotFoundException" | ||
|
|
||
| @pytest.mark.order(6) | ||
| def test_dict_and_string_equivalence(self): | ||
| """put(dict) and put(json.dumps(dict)) produce the same get() result.""" | ||
| policy = _make_policy(self.resource_arn, self.principal_arn) | ||
|
|
||
| self.client.put_resource_policy(self.resource_arn, policy) | ||
| result_from_dict = self.client.get_resource_policy(self.resource_arn) | ||
|
|
||
| self.client.put_resource_policy(self.resource_arn, json.dumps(policy)) | ||
| result_from_str = self.client.get_resource_policy(self.resource_arn) | ||
|
|
||
| assert result_from_dict == result_from_str |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Managing an agent runtime requires two clients: the existing runtime client and a resource policy client. Customers who need to attach resource policies to a runtime will need this additional client.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is the tradeoff with this approach. If we hide the resource policy functionality behind other clients, the customers can manage an agent runtime with a single client, but we are responsible for wiring that up each time we want to add support for a new primitive or agentcore resource.
Alternatively, if we create a separate class the customer will need separate clients for the mentioned flow, but it allows for a more extendable and flexible design that doesn't require explicitly wiring up each consumer.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I understand why having a separate client. Thank you for weighing both options. I'm fine with merging this client.