From ecf502e9b912c538e3fec909611e80b81fdce7cf Mon Sep 17 00:00:00 2001 From: Nicolas Borges Date: Mon, 20 Apr 2026 12:56:01 -0400 Subject: [PATCH 1/2] feat: update runtime client with passthrough --- .../runtime/agent_core_runtime_client.py | 335 +++++++++++++++++- .../test_region_validation.py | 4 +- .../unit/runtime/test_runtime_passthrough.py | 208 +++++++++++ .../runtime/test_runtime_passthrough.py | 174 +++++++++ 4 files changed, 708 insertions(+), 13 deletions(-) create mode 100644 tests/unit/runtime/test_runtime_passthrough.py create mode 100644 tests_integ/runtime/test_runtime_passthrough.py diff --git a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py b/src/bedrock_agentcore/runtime/agent_core_runtime_client.py index ab91de5a..8717e9b9 100644 --- a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py +++ b/src/bedrock_agentcore/runtime/agent_core_runtime_client.py @@ -9,19 +9,28 @@ import logging import secrets import uuid -from typing import Dict, Optional, Tuple +from typing import Any, Dict, Optional, Tuple from urllib.parse import quote, urlencode, urlparse import boto3 from botocore.auth import SigV4Auth, SigV4QueryAuth from botocore.awsrequest import AWSRequest +from botocore.config import Config +from botocore.exceptions import ClientError +from .._utils.config import WaitConfig from .._utils.endpoints import get_data_plane_endpoint +from .._utils.polling import wait_until, wait_until_deleted +from .._utils.snake_case import accept_snake_case_kwargs, convert_kwargs +from .._utils.user_agent import build_user_agent_suffix from .utils import is_valid_partition DEFAULT_PRESIGNED_URL_TIMEOUT = 300 MAX_PRESIGNED_URL_TIMEOUT = 300 +_RUNTIME_FAILED_STATUSES = {"CREATE_FAILED", "UPDATE_FAILED"} +_ENDPOINT_FAILED_STATUSES = {"CREATE_FAILED", "UPDATE_FAILED", "DELETE_FAILED"} + class AgentCoreRuntimeClient: """Client for generating WebSocket authentication for AgentCore Runtime. @@ -35,24 +44,86 @@ class AgentCoreRuntimeClient: session (boto3.Session): The boto3 session for AWS credentials. """ - def __init__(self, region: str, session: Optional[boto3.Session] = None) -> None: + _ALLOWED_DP_METHODS = { + "invoke_agent_runtime", + "stop_runtime_session", + } + + _ALLOWED_CP_METHODS = { + "create_agent_runtime", + "update_agent_runtime", + "get_agent_runtime", + "delete_agent_runtime", + "list_agent_runtimes", + "create_agent_runtime_endpoint", + "get_agent_runtime_endpoint", + "update_agent_runtime_endpoint", + "delete_agent_runtime_endpoint", + "list_agent_runtime_endpoints", + "list_agent_runtime_versions", + "delete_agent_runtime_version", + } + + def __init__( + self, + region: Optional[str] = None, + session: Optional[boto3.Session] = None, + integration_source: Optional[str] = None, + ) -> None: """Initialize an AgentCoreRuntime client for the specified AWS region. Args: - region (str): The AWS region to use for the AgentCore Runtime service. - session (Optional[boto3.Session]): Optional boto3 session. If not provided, - a new session will be created using default credentials. + region: AWS region name. If not provided, uses the session's + region or "us-west-2". + session: Optional boto3 Session to use. If not provided, a + default session is created. + integration_source: Optional integration source for user-agent + telemetry. """ - from .._utils.endpoints import validate_region - - validate_region(region) - self.region = region + session = session if session else boto3.Session() + self.region = region or session.region_name or "us-west-2" + self.session = session + self.integration_source = integration_source self.logger = logging.getLogger(__name__) - if session is None: - session = boto3.Session() + user_agent_extra = build_user_agent_suffix(integration_source) + client_config = Config(user_agent_extra=user_agent_extra) - self.session = session + self.cp_client = session.client( + "bedrock-agentcore-control", + region_name=self.region, + config=client_config, + ) + self.dp_client = session.client( + "bedrock-agentcore", + region_name=self.region, + config=client_config, + ) + self.logger.info( + "Initialized AgentCoreRuntimeClient for region: %s", + self.region, + ) + + # Pass-through + # ------------------------------------------------------------------------- + def __getattr__(self, name: str): + """Dynamically forward allowlisted method calls to the appropriate boto3 client.""" + if name in self._ALLOWED_DP_METHODS and hasattr(self.dp_client, name): + method = getattr(self.dp_client, name) + self.logger.debug("Forwarding method '%s' to dp_client", name) + return accept_snake_case_kwargs(method) + + if name in self._ALLOWED_CP_METHODS and hasattr(self.cp_client, name): + method = getattr(self.cp_client, name) + self.logger.debug("Forwarding method '%s' to cp_client", name) + return accept_snake_case_kwargs(method) + + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{name}'. " + f"Method not found on dp_client or cp_client. " + f"Available methods can be found in the boto3 documentation for " + f"'bedrock-agentcore' and 'bedrock-agentcore-control' services." + ) def _parse_runtime_arn(self, runtime_arn: str) -> Dict[str, str]: """Parse runtime ARN and extract components. @@ -401,3 +472,243 @@ def generate_ws_connection_oauth( self.logger.debug("Bearer token length: %d characters", len(bearer_token)) return ws_url, headers + + # *_and_wait methods + # ------------------------------------------------------------------------- + def create_agent_runtime_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Create an agent runtime and wait for it to reach READY status. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the create_agent_runtime API. + + Returns: + Runtime details when READY. + + Raises: + RuntimeError: If the runtime reaches a failed state. + TimeoutError: If the runtime doesn't become READY within max_wait. + """ + response = self.cp_client.create_agent_runtime(**convert_kwargs(kwargs)) + rid = response["agentRuntimeId"] + return wait_until( + lambda: self.cp_client.get_agent_runtime(agentRuntimeId=rid), + "READY", + _RUNTIME_FAILED_STATUSES, + wait_config, + error_field="failureReason", + ) + + def update_agent_runtime_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Update an agent runtime and wait for it to reach READY status. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the update_agent_runtime API. + + Returns: + Runtime details when READY. + + Raises: + RuntimeError: If the runtime reaches a failed state. + TimeoutError: If the runtime doesn't become READY within max_wait. + """ + response = self.cp_client.update_agent_runtime(**convert_kwargs(kwargs)) + rid = response["agentRuntimeId"] + return wait_until( + lambda: self.cp_client.get_agent_runtime(agentRuntimeId=rid), + "READY", + _RUNTIME_FAILED_STATUSES, + wait_config, + error_field="failureReason", + ) + + def delete_agent_runtime_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> None: + """Delete an agent runtime and wait for deletion to complete. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the delete_agent_runtime API. + + Raises: + TimeoutError: If the runtime isn't deleted within max_wait. + """ + response = self.cp_client.delete_agent_runtime(**convert_kwargs(kwargs)) + rid = response["agentRuntimeId"] + wait_until_deleted( + lambda: self.cp_client.get_agent_runtime(agentRuntimeId=rid), + wait_config=wait_config, + ) + + def create_agent_runtime_endpoint_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Create an agent runtime endpoint and wait for it to reach READY. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the + create_agent_runtime_endpoint API. + + Returns: + Endpoint details when READY. + + Raises: + RuntimeError: If the endpoint reaches a failed state. + TimeoutError: If the endpoint doesn't become READY within + max_wait. + """ + response = self.cp_client.create_agent_runtime_endpoint( + **convert_kwargs(kwargs), + ) + rid = kwargs.get( + "agentRuntimeId", + convert_kwargs(kwargs).get("agentRuntimeId"), + ) + ename = response.get("name", kwargs.get("name", "DEFAULT")) + return wait_until( + lambda: self.cp_client.get_agent_runtime_endpoint( + agentRuntimeId=rid, + endpointName=ename, + ), + "READY", + _ENDPOINT_FAILED_STATUSES, + wait_config, + error_field="failureReason", + ) + + def update_agent_runtime_endpoint_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Update an agent runtime endpoint and wait for READY status. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the + update_agent_runtime_endpoint API. + + Returns: + Endpoint details when READY. + + Raises: + RuntimeError: If the endpoint reaches a failed state. + TimeoutError: If the endpoint doesn't become READY within + max_wait. + """ + response = self.cp_client.update_agent_runtime_endpoint( + **convert_kwargs(kwargs), + ) + rid = kwargs.get( + "agentRuntimeId", + convert_kwargs(kwargs).get("agentRuntimeId"), + ) + ename = response.get("name", kwargs.get("endpointName", "DEFAULT")) + return wait_until( + lambda: self.cp_client.get_agent_runtime_endpoint( + agentRuntimeId=rid, + endpointName=ename, + ), + "READY", + _ENDPOINT_FAILED_STATUSES, + wait_config, + error_field="failureReason", + ) + + # Higher-level orchestration methods + # ------------------------------------------------------------------------- + def get_aggregated_status( + self, + agent_runtime_id: str, + endpoint_name: str = "DEFAULT", + ) -> Dict[str, Any]: + """Get aggregated status of runtime and endpoint. + + Args: + agent_runtime_id: The agent runtime ID. + endpoint_name: Endpoint name (default: "DEFAULT"). + + Returns: + Dict with 'runtime' and 'endpoint' status details. + """ + result: Dict[str, Any] = {"runtime": None, "endpoint": None} + + try: + result["runtime"] = self.cp_client.get_agent_runtime( + agentRuntimeId=agent_runtime_id, + ) + except ClientError as e: + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise + result["runtime"] = {"error": str(e)} + + try: + result["endpoint"] = self.cp_client.get_agent_runtime_endpoint( + agentRuntimeId=agent_runtime_id, + endpointName=endpoint_name, + ) + except ClientError as e: + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise + result["endpoint"] = {"error": str(e)} + + return result + + def teardown_endpoint_and_runtime( + self, + agent_runtime_id: str, + endpoint_name: str = "DEFAULT", + ) -> None: + """Delete endpoint then runtime in correct order. + + Silently ignores ResourceNotFoundException for either resource + (already deleted). + + Args: + agent_runtime_id: The agent runtime ID. + endpoint_name: Endpoint name (default: "DEFAULT"). + """ + try: + self.cp_client.delete_agent_runtime_endpoint( + agentRuntimeId=agent_runtime_id, + endpointName=endpoint_name, + ) + self.logger.info( + "Deleted endpoint '%s' for runtime %s", + endpoint_name, + agent_runtime_id, + ) + wait_until_deleted( + lambda: self.cp_client.get_agent_runtime_endpoint( + agentRuntimeId=agent_runtime_id, + endpointName=endpoint_name, + ), + ) + except ClientError as e: + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise + self.logger.info("Endpoint '%s' not found, skipping", endpoint_name) + + try: + self.delete_agent_runtime_and_wait( + agentRuntimeId=agent_runtime_id, + ) + except ClientError as e: + if e.response["Error"]["Code"] != "ResourceNotFoundException": + raise + self.logger.info("Runtime %s not found, skipping", agent_runtime_id) diff --git a/tests/bedrock_agentcore/test_region_validation.py b/tests/bedrock_agentcore/test_region_validation.py index cee328f7..9f942c20 100644 --- a/tests/bedrock_agentcore/test_region_validation.py +++ b/tests/bedrock_agentcore/test_region_validation.py @@ -219,7 +219,9 @@ class TestClientConstructorValidation: """Tests that all client constructors reject malicious regions early.""" def test_agent_core_runtime_client_rejects_bad_region(self): - with pytest.raises(InvalidRegionError): + from botocore.exceptions import InvalidRegionError as BotocoreInvalidRegionError + + with pytest.raises(BotocoreInvalidRegionError): from bedrock_agentcore.runtime.agent_core_runtime_client import AgentCoreRuntimeClient AgentCoreRuntimeClient("x@attacker.com:443/#") diff --git a/tests/unit/runtime/test_runtime_passthrough.py b/tests/unit/runtime/test_runtime_passthrough.py new file mode 100644 index 00000000..6aea413d --- /dev/null +++ b/tests/unit/runtime/test_runtime_passthrough.py @@ -0,0 +1,208 @@ +"""Tests for AgentCoreRuntimeClient passthrough, *_and_wait, and orchestration methods.""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest +from botocore.exceptions import ClientError + +from bedrock_agentcore.runtime.agent_core_runtime_client import AgentCoreRuntimeClient + + +class TestRuntimeClientPassthrough: + """Tests for __getattr__ passthrough.""" + + def _make_client(self): + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + client = AgentCoreRuntimeClient(session=mock_session) + client.cp_client = Mock() + client.dp_client = Mock() + return client + + def test_cp_method_forwarded(self): + client = self._make_client() + client.cp_client.get_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + result = client.get_agent_runtime(agentRuntimeId="r-123") + client.cp_client.get_agent_runtime.assert_called_once_with(agentRuntimeId="r-123") + assert result["agentRuntimeId"] == "r-123" + + def test_dp_method_forwarded(self): + client = self._make_client() + client.dp_client.invoke_agent_runtime.return_value = {"output": "hello"} + result = client.invoke_agent_runtime(agentRuntimeId="r-123") + client.dp_client.invoke_agent_runtime.assert_called_once_with(agentRuntimeId="r-123") + assert result["output"] == "hello" + + def test_snake_case_kwargs_converted(self): + client = self._make_client() + client.cp_client.get_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.get_agent_runtime(agent_runtime_id="r-123") + client.cp_client.get_agent_runtime.assert_called_once_with(agentRuntimeId="r-123") + + def test_non_allowlisted_method_raises(self): + client = self._make_client() + with pytest.raises(AttributeError, match="has no attribute"): + client.not_a_real_method() + + def test_all_cp_methods_in_allowlist(self): + expected = { + "create_agent_runtime", + "update_agent_runtime", + "get_agent_runtime", + "delete_agent_runtime", + "list_agent_runtimes", + "create_agent_runtime_endpoint", + "get_agent_runtime_endpoint", + "update_agent_runtime_endpoint", + "delete_agent_runtime_endpoint", + "list_agent_runtime_endpoints", + "list_agent_runtime_versions", + "delete_agent_runtime_version", + } + assert expected == AgentCoreRuntimeClient._ALLOWED_CP_METHODS + + def test_all_dp_methods_in_allowlist(self): + expected = {"invoke_agent_runtime", "stop_runtime_session"} + assert expected == AgentCoreRuntimeClient._ALLOWED_DP_METHODS + + +class TestRuntimeAndWait: + """Tests for *_and_wait methods.""" + + def _make_client(self): + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + client = AgentCoreRuntimeClient(session=mock_session) + client.cp_client = Mock() + client.dp_client = Mock() + return client + + def test_create_agent_runtime_and_wait(self): + client = self._make_client() + client.cp_client.create_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.return_value = { + "status": "READY", + "agentRuntimeId": "r-123", + } + result = client.create_agent_runtime_and_wait(agentRuntimeName="test") + assert result["status"] == "READY" + + def test_create_agent_runtime_and_wait_failed(self): + client = self._make_client() + client.cp_client.create_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.return_value = { + "status": "CREATE_FAILED", + "failureReason": "bad config", + } + with pytest.raises(RuntimeError, match="CREATE_FAILED"): + client.create_agent_runtime_and_wait(agentRuntimeName="test") + + def test_update_agent_runtime_and_wait(self): + client = self._make_client() + client.cp_client.update_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.return_value = { + "status": "READY", + "agentRuntimeId": "r-123", + } + result = client.update_agent_runtime_and_wait(agentRuntimeId="r-123") + assert result["status"] == "READY" + + @patch("bedrock_agentcore._utils.polling.time.sleep") + def test_delete_agent_runtime_and_wait(self, _mock_sleep): + client = self._make_client() + client.cp_client.delete_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": "gone"}}, + "GetAgentRuntime", + ) + client.delete_agent_runtime_and_wait(agentRuntimeId="r-123") + client.cp_client.delete_agent_runtime.assert_called_once() + + def test_get_agent_runtime_endpoint_and_wait(self): + client = self._make_client() + client.cp_client.update_agent_runtime_endpoint.return_value = { + "name": "DEFAULT", + } + client.cp_client.get_agent_runtime_endpoint.return_value = { + "status": "READY", + "endpointName": "DEFAULT", + } + result = client.update_agent_runtime_endpoint_and_wait( + agentRuntimeId="r-123", + endpointName="DEFAULT", + ) + assert result["status"] == "READY" + + +class TestRuntimeOrchestration: + """Tests for get_aggregated_status and teardown.""" + + def _make_client(self): + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + client = AgentCoreRuntimeClient(session=mock_session) + client.cp_client = Mock() + client.dp_client = Mock() + return client + + def test_get_aggregated_status(self): + client = self._make_client() + client.cp_client.get_agent_runtime.return_value = {"status": "READY"} + client.cp_client.get_agent_runtime_endpoint.return_value = {"status": "READY"} + result = client.get_aggregated_status("r-123") + assert result["runtime"]["status"] == "READY" + assert result["endpoint"]["status"] == "READY" + + def test_get_aggregated_status_not_found(self): + client = self._make_client() + client.cp_client.get_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "GetAgentRuntime", + ) + client.cp_client.get_agent_runtime_endpoint.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "GetAgentRuntimeEndpoint", + ) + result = client.get_aggregated_status("r-123") + assert "error" in result["runtime"] + assert "error" in result["endpoint"] + + def test_get_aggregated_status_reraises_non_not_found(self): + client = self._make_client() + client.cp_client.get_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "denied"}}, + "GetAgentRuntime", + ) + with pytest.raises(ClientError): + client.get_aggregated_status("r-123") + + @patch("bedrock_agentcore._utils.polling.time.sleep") + def test_teardown_deletes_endpoint_then_runtime(self, _mock_sleep): + client = self._make_client() + client.cp_client.get_agent_runtime_endpoint.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "GetAgentRuntimeEndpoint", + ) + client.cp_client.delete_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "GetAgentRuntime", + ) + client.teardown_endpoint_and_runtime("r-123") + client.cp_client.delete_agent_runtime_endpoint.assert_called_once() + client.cp_client.delete_agent_runtime.assert_called_once() + + @patch("bedrock_agentcore._utils.polling.time.sleep") + def test_teardown_skips_missing_endpoint(self, _mock_sleep): + client = self._make_client() + client.cp_client.delete_agent_runtime_endpoint.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "DeleteAgentRuntimeEndpoint", + ) + client.cp_client.delete_agent_runtime.return_value = {"agentRuntimeId": "r-123"} + client.cp_client.get_agent_runtime.side_effect = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": ""}}, + "GetAgentRuntime", + ) + client.teardown_endpoint_and_runtime("r-123") + client.cp_client.delete_agent_runtime.assert_called_once() diff --git a/tests_integ/runtime/test_runtime_passthrough.py b/tests_integ/runtime/test_runtime_passthrough.py new file mode 100644 index 00000000..23ad2272 --- /dev/null +++ b/tests_integ/runtime/test_runtime_passthrough.py @@ -0,0 +1,174 @@ +"""Integration tests for AgentCoreRuntimeClient passthrough and *_and_wait methods. + +Requires environment variables: + BEDROCK_TEST_REGION: AWS region (default: us-west-2) + RUNTIME_ROLE_ARN: IAM role ARN with AgentCore runtime trust policy + RUNTIME_S3_CODE_URI: S3 URI for agent code artifact (e.g. s3://bucket/agent.zip) +""" + +import os +import time + +import pytest +from botocore.exceptions import ClientError + +from bedrock_agentcore.runtime.agent_core_runtime_client import AgentCoreRuntimeClient + + +@pytest.mark.integration +class TestRuntimeClientPassthrough: + """Read-only passthrough tests. No resources needed.""" + + @classmethod + def setup_class(cls): + cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-west-2") + cls.client = AgentCoreRuntimeClient(region=cls.region) + + @pytest.mark.order(1) + def test_list_agent_runtimes_passthrough(self): + response = self.client.list_agent_runtimes() + assert "agentRuntimes" in response + + @pytest.mark.order(2) + def test_list_agent_runtimes_snake_case(self): + response = self.client.list_agent_runtimes(max_results=5) + assert "agentRuntimes" in response + + @pytest.mark.order(3) + def test_non_allowlisted_method_raises(self): + with pytest.raises(AttributeError): + self.client.not_a_real_method() + + +@pytest.mark.integration +class TestRuntimeCrud: + """CRUD tests for agent runtimes. + + Requires RUNTIME_ROLE_ARN and RUNTIME_S3_CODE_URI. + """ + + @classmethod + def setup_class(cls): + cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-west-2") + cls.role_arn = os.environ.get("RUNTIME_ROLE_ARN") + cls.s3_code_uri = os.environ.get("RUNTIME_S3_CODE_URI") + if not cls.role_arn or not cls.s3_code_uri: + pytest.skip("RUNTIME_ROLE_ARN and RUNTIME_S3_CODE_URI must be set") + cls.client = AgentCoreRuntimeClient(region=cls.region) + cls.test_prefix = f"sdk_integ_{int(time.time())}" + cls.runtime_ids = [] + + # Parse S3 URI into bucket and key + s3_parts = cls.s3_code_uri.replace("s3://", "").split("/", 1) + cls.s3_bucket = s3_parts[0] + cls.s3_key = s3_parts[1] + + @classmethod + def teardown_class(cls): + for rid in cls.runtime_ids: + try: + cls.client.teardown_endpoint_and_runtime( + rid, + endpoint_name="sdk_integ_ep", + ) + except Exception as e: + print(f"Failed to teardown runtime {rid}: {e}") + + @pytest.mark.order(4) + def test_create_agent_runtime_and_wait(self): + runtime = self.client.create_agent_runtime_and_wait( + agentRuntimeName=f"{self.test_prefix}_rt", + roleArn=self.role_arn, + networkConfiguration={"networkMode": "PUBLIC"}, + agentRuntimeArtifact={ + "codeConfiguration": { + "code": { + "s3": { + "bucket": self.s3_bucket, + "prefix": self.s3_key, + } + }, + "runtime": "PYTHON_3_12", + "entryPoint": ["agent.py"], + } + }, + ) + self.__class__.runtime_ids.append(runtime["agentRuntimeId"]) + assert runtime["status"] == "READY" + + @pytest.mark.order(5) + def test_get_agent_runtime_passthrough(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + runtime = self.client.get_agent_runtime( + agentRuntimeId=self.runtime_ids[0], + ) + assert runtime["status"] == "READY" + + @pytest.mark.order(6) + def test_get_agent_runtime_snake_case(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + runtime = self.client.get_agent_runtime( + agent_runtime_id=self.runtime_ids[0], + ) + assert runtime["status"] == "READY" + + @pytest.mark.order(7) + def test_update_agent_runtime_and_wait(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + updated = self.client.update_agent_runtime_and_wait( + agentRuntimeId=self.runtime_ids[0], + roleArn=self.role_arn, + networkConfiguration={"networkMode": "PUBLIC"}, + agentRuntimeArtifact={ + "codeConfiguration": { + "code": { + "s3": { + "bucket": self.s3_bucket, + "prefix": self.s3_key, + } + }, + "runtime": "PYTHON_3_12", + "entryPoint": ["agent.py"], + } + }, + description="updated by integ test", + ) + assert updated["status"] == "READY" + + @pytest.mark.order(8) + def test_create_agent_runtime_endpoint_and_wait(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + endpoint = self.client.create_agent_runtime_endpoint_and_wait( + agentRuntimeId=self.runtime_ids[0], + name="sdk_integ_ep", + ) + assert endpoint["status"] == "READY" + + @pytest.mark.order(9) + def test_get_aggregated_status(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + result = self.client.get_aggregated_status( + self.runtime_ids[0], + endpoint_name="sdk_integ_ep", + ) + assert result["runtime"] is not None + assert result["runtime"]["status"] == "READY" + assert result["endpoint"] is not None + assert result["endpoint"]["status"] == "READY" + + @pytest.mark.order(10) + def test_teardown_endpoint_and_runtime(self): + if not self.runtime_ids: + pytest.skip("prerequisite test did not create runtime") + rid = self.runtime_ids.pop(0) + self.client.teardown_endpoint_and_runtime( + rid, + endpoint_name="sdk_integ_ep", + ) + with pytest.raises(ClientError): + self.client.get_agent_runtime(agentRuntimeId=rid) From 00cb54fb1b0dbd0494122d273237c9c75fdb1581 Mon Sep 17 00:00:00 2001 From: Nicolas Borges Date: Fri, 24 Apr 2026 13:27:56 -0400 Subject: [PATCH 2/2] fix: validate region on construction instead of methods --- .../runtime/agent_core_runtime_client.py | 20 ++++++++----------- .../test_region_validation.py | 4 +--- 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py b/src/bedrock_agentcore/runtime/agent_core_runtime_client.py index 8717e9b9..27c699aa 100644 --- a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py +++ b/src/bedrock_agentcore/runtime/agent_core_runtime_client.py @@ -19,7 +19,7 @@ from botocore.exceptions import ClientError from .._utils.config import WaitConfig -from .._utils.endpoints import get_data_plane_endpoint +from .._utils.endpoints import get_data_plane_endpoint, validate_region from .._utils.polling import wait_until, wait_until_deleted from .._utils.snake_case import accept_snake_case_kwargs, convert_kwargs from .._utils.user_agent import build_user_agent_suffix @@ -81,7 +81,7 @@ def __init__( telemetry. """ session = session if session else boto3.Session() - self.region = region or session.region_name or "us-west-2" + self.region = validate_region(region or session.region_name or "us-west-2") self.session = session self.integration_source = integration_source self.logger = logging.getLogger(__name__) @@ -572,13 +572,11 @@ def create_agent_runtime_endpoint_and_wait( TimeoutError: If the endpoint doesn't become READY within max_wait. """ + converted = convert_kwargs(kwargs) response = self.cp_client.create_agent_runtime_endpoint( - **convert_kwargs(kwargs), - ) - rid = kwargs.get( - "agentRuntimeId", - convert_kwargs(kwargs).get("agentRuntimeId"), + **converted, ) + rid = converted.get("agentRuntimeId") ename = response.get("name", kwargs.get("name", "DEFAULT")) return wait_until( lambda: self.cp_client.get_agent_runtime_endpoint( @@ -611,13 +609,11 @@ def update_agent_runtime_endpoint_and_wait( TimeoutError: If the endpoint doesn't become READY within max_wait. """ + converted = convert_kwargs(kwargs) response = self.cp_client.update_agent_runtime_endpoint( - **convert_kwargs(kwargs), - ) - rid = kwargs.get( - "agentRuntimeId", - convert_kwargs(kwargs).get("agentRuntimeId"), + **converted, ) + rid = converted.get("agentRuntimeId") ename = response.get("name", kwargs.get("endpointName", "DEFAULT")) return wait_until( lambda: self.cp_client.get_agent_runtime_endpoint( diff --git a/tests/bedrock_agentcore/test_region_validation.py b/tests/bedrock_agentcore/test_region_validation.py index 9f942c20..cee328f7 100644 --- a/tests/bedrock_agentcore/test_region_validation.py +++ b/tests/bedrock_agentcore/test_region_validation.py @@ -219,9 +219,7 @@ class TestClientConstructorValidation: """Tests that all client constructors reject malicious regions early.""" def test_agent_core_runtime_client_rejects_bad_region(self): - from botocore.exceptions import InvalidRegionError as BotocoreInvalidRegionError - - with pytest.raises(BotocoreInvalidRegionError): + with pytest.raises(InvalidRegionError): from bedrock_agentcore.runtime.agent_core_runtime_client import AgentCoreRuntimeClient AgentCoreRuntimeClient("x@attacker.com:443/#")