From 67424f417dd5e011715222fa9c96e5b9e3d42d1d Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Sun, 29 Mar 2026 15:19:15 -0700 Subject: [PATCH] fix: lint errors --- .../tools/code_interpreter_client.py | 15 +++ .../tools/test_code_interpreter_client.py | 124 ++++++++++++++++++ tests_integ/tools/test_code.py | 96 +++++++++++++- 3 files changed, 234 insertions(+), 1 deletion(-) diff --git a/src/bedrock_agentcore/tools/code_interpreter_client.py b/src/bedrock_agentcore/tools/code_interpreter_client.py index 636b7eb5..7c2ab4f6 100644 --- a/src/bedrock_agentcore/tools/code_interpreter_client.py +++ b/src/bedrock_agentcore/tools/code_interpreter_client.py @@ -16,10 +16,19 @@ from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint, get_data_plane_endpoint from bedrock_agentcore._utils.user_agent import build_user_agent_suffix +from .config import Certificate + DEFAULT_IDENTIFIER = "aws.codeinterpreter.v1" DEFAULT_TIMEOUT = 900 +def _to_dict(obj): + """Convert an object to a dict, calling to_dict() if available.""" + if hasattr(obj, "to_dict"): + return obj.to_dict() + return obj + + class CodeInterpreter: """Client for interacting with the AWS Code Interpreter sandbox service. @@ -134,6 +143,7 @@ def create_code_interpreter( execution_role_arn: str, network_configuration: Optional[Dict] = None, description: Optional[str] = None, + certificates: Optional[List[Union[Certificate, Dict[str, Any]]]] = None, tags: Optional[Dict[str, str]] = None, client_token: Optional[str] = None, ) -> Dict: @@ -155,6 +165,8 @@ def create_code_interpreter( } } description (Optional[str]): Description of the interpreter (1-4096 chars) + certificates (Optional[List[Union[Certificate, Dict]]]): Root CA certificates + from Secrets Manager for the code interpreter to trust. tags (Optional[Dict[str, str]]): Tags for the interpreter client_token (Optional[str]): Idempotency token @@ -193,6 +205,9 @@ def create_code_interpreter( if description: request_params["description"] = description + if certificates: + request_params["certificates"] = [_to_dict(c) for c in certificates] + if tags: request_params["tags"] = tags diff --git a/tests/bedrock_agentcore/tools/test_code_interpreter_client.py b/tests/bedrock_agentcore/tools/test_code_interpreter_client.py index 59fcead0..9d039b58 100644 --- a/tests/bedrock_agentcore/tools/test_code_interpreter_client.py +++ b/tests/bedrock_agentcore/tools/test_code_interpreter_client.py @@ -1500,3 +1500,127 @@ def test_execute_code_auto_starts_session(self, mock_boto3, mock_get_data_endpoi # Assert client.data_plane_client.start_code_interpreter_session.assert_called_once() client.data_plane_client.invoke_code_interpreter.assert_called_once() + + @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") + def test_create_code_interpreter_with_certificates( + self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint + ): + # Arrange + mock_session = MagicMock() + mock_control_client = MagicMock() + mock_data_client = MagicMock() + mock_session.client.side_effect = [mock_control_client, mock_data_client] + mock_boto3.Session.return_value = mock_session + client = CodeInterpreter("us-west-2") + + mock_response = { + "codeInterpreterArn": "arn:aws:bedrock-agentcore:us-west-2:123456789012:code-interpreter/test-interp", + "codeInterpreterId": "test-interp-123", + "createdAt": datetime.datetime.now(), + "status": "CREATING", + } + client.control_plane_client.create_code_interpreter.return_value = mock_response + + certificates = [ + { + "location": { + "secretsManager": {"secretArn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:my-ca"} + } + } + ] + + # Act + result = client.create_code_interpreter( + name="test_interpreter_with_certs", + execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole", + certificates=certificates, + description="Interpreter with custom root CA", + ) + + # Assert + client.control_plane_client.create_code_interpreter.assert_called_once_with( + name="test_interpreter_with_certs", + executionRoleArn="arn:aws:iam::123456789012:role/InterpreterRole", + networkConfiguration={"networkMode": "PUBLIC"}, + description="Interpreter with custom root CA", + certificates=certificates, + ) + assert result == mock_response + + @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") + def test_create_code_interpreter_with_certificate_dataclass( + self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint + ): + # Arrange + mock_session = MagicMock() + mock_control_client = MagicMock() + mock_data_client = MagicMock() + mock_session.client.side_effect = [mock_control_client, mock_data_client] + mock_boto3.Session.return_value = mock_session + client = CodeInterpreter("us-west-2") + + mock_response = { + "codeInterpreterArn": "arn:aws:bedrock-agentcore:us-west-2:123456789012:code-interpreter/test-interp", + "codeInterpreterId": "test-interp-123", + "createdAt": datetime.datetime.now(), + "status": "CREATING", + } + client.control_plane_client.create_code_interpreter.return_value = mock_response + + # Use a mock Certificate dataclass with to_dict method + mock_cert = MagicMock() + mock_cert.to_dict.return_value = { + "location": {"secretsManager": {"secretArn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:my-ca"}} + } + + # Act + result = client.create_code_interpreter( + name="test_interpreter_with_certs", + execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole", + certificates=[mock_cert], + ) + + # Assert + mock_cert.to_dict.assert_called_once() + client.control_plane_client.create_code_interpreter.assert_called_once_with( + name="test_interpreter_with_certs", + executionRoleArn="arn:aws:iam::123456789012:role/InterpreterRole", + networkConfiguration={"networkMode": "PUBLIC"}, + certificates=[mock_cert.to_dict.return_value], + ) + assert result == mock_response + + @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint") + @patch("bedrock_agentcore.tools.code_interpreter_client.boto3") + def test_create_code_interpreter_without_certificates( + self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint + ): + """Verify that omitting certificates does not include it in the API call (backward compatibility).""" + # Arrange + mock_session = MagicMock() + mock_control_client = MagicMock() + mock_data_client = MagicMock() + mock_session.client.side_effect = [mock_control_client, mock_data_client] + mock_boto3.Session.return_value = mock_session + client = CodeInterpreter("us-west-2") + + mock_response = { + "codeInterpreterId": "test-interp-123", + "status": "CREATING", + } + client.control_plane_client.create_code_interpreter.return_value = mock_response + + # Act + client.create_code_interpreter( + name="test_interpreter_no_certs", + execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole", + ) + + # Assert — certificates key should NOT be in the call + call_kwargs = client.control_plane_client.create_code_interpreter.call_args[1] + assert "certificates" not in call_kwargs diff --git a/tests_integ/tools/test_code.py b/tests_integ/tools/test_code.py index 1b02c95a..070b209f 100644 --- a/tests_integ/tools/test_code.py +++ b/tests_integ/tools/test_code.py @@ -4,7 +4,13 @@ To run: pytest tests_integ/tools/test_code.py -v """ -from bedrock_agentcore.tools.code_interpreter_client import code_session +import os +import time + +import boto3 + +from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint +from bedrock_agentcore.tools.code_interpreter_client import CodeInterpreter, code_session # Test 1: Basic code execution with system interpreter print("Test 1: Basic code execution using execute_code()") @@ -223,6 +229,94 @@ print(event["result"]["content"]) print("✅ Test 10 passed\n") +# Test 11: Create code interpreter with custom root CA certificate +print("Test 11: Code interpreter with custom root CA certificate") + +REGION = "us-west-2" +secret_arn = os.environ.get("TEST_ROOT_CA_SECRET_ARN") +execution_role_arn = os.environ.get("TEST_EXECUTION_ROLE_ARN") + +if secret_arn and execution_role_arn: + cp_client = boto3.client( + "bedrock-agentcore-control", + region_name=REGION, + endpoint_url=get_control_plane_endpoint(REGION), + ) + + # Create interpreter with root CA + import time + + response = cp_client.create_code_interpreter( + name=f"integ_test_rootca_{int(time.time())}", + executionRoleArn=execution_role_arn, + networkConfiguration={"networkMode": "PUBLIC"}, + certificates=[{"location": {"secretsManager": {"secretArn": secret_arn}}}], + description="Integration test: code interpreter with custom root CA", + ) + interpreter_id = response["codeInterpreterId"] + print(f" Created interpreter: {interpreter_id}") + + # Wait for ready + ci_client = CodeInterpreter(REGION) + for _ in range(40): + info = ci_client.get_code_interpreter(interpreter_id) + if info["status"] == "READY": + break + elif info["status"] == "CREATE_FAILED": + raise RuntimeError(f"Interpreter creation failed: {info.get('failureReason')}") + time.sleep(3) + + print(f" Interpreter status: {info['status']}") + + # Start session and test SSL connection to untrusted-root.badssl.com + ci_client.start(identifier=interpreter_id) + print(f" Session started: {ci_client.session_id}") + + result = ci_client.invoke( + "executeCode", + { + "code": ( + "import urllib.request\n" + "response = urllib.request.urlopen(" + '"https://untrusted-root.badssl.com")\n' + 'print(f"Status: {response.status}")' + ), + "language": "python", + }, + ) + + success = False + for event in result.get("stream", []): + if "result" in event: + structured = event["result"].get("structuredContent", {}) + stdout = structured.get("stdout", "") + if "200" in stdout: + success = True + print(f" SSL connection succeeded: {stdout.strip()}") + + ci_client.stop() + + # Cleanup - delete interpreter + sessions = ci_client.list_sessions(interpreter_id=interpreter_id, status="READY") + for s in sessions.get("items", []): + try: + ci_client.data_plane_client.stop_code_interpreter_session( + codeInterpreterIdentifier=interpreter_id, sessionId=s["sessionId"] + ) + except Exception: + pass + time.sleep(5) + cp_client.delete_code_interpreter(codeInterpreterId=interpreter_id) + print(f" Cleaned up interpreter: {interpreter_id}") + + assert success, "SSL connection to untrusted-root.badssl.com should succeed with custom root CA" + print("✅ Test 11 passed\n") +else: + print(" ⏭️ Skipped (set TEST_ROOT_CA_SECRET_ARN and TEST_EXECUTION_ROLE_ARN env vars to enable)") + print(" Example:") + print(" export TEST_ROOT_CA_SECRET_ARN=arn:aws:secretsmanager:us-west-2:123456789012:secret:badssl-root-ca") + print(" export TEST_EXECUTION_ROLE_ARN=arn:aws:iam::123456789012:role/AgentCoreRole") + print("✅ Test 11 skipped\n") print("=" * 50) print("All integration tests passed! ✅")