fix(deps): restrict pydantic to versions below 2.41.3 (#280)

* fix(deps): restrict pydantic to versions below 2.41.3

* fix integration test to use auto create for execution role

---------

Co-authored-by: Sundar Raghavan <sdraghav@amazon.com>

Author: sundargthb

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "httpx>=0.28.1",
     "jinja2>=3.1.6",
     "prompt-toolkit>=3.0.51",
-    "pydantic>=2.0.0,<3.0.0",
+    "pydantic>=2.0.0,<2.41.3",
     "urllib3>=1.26.0",
     "pyyaml>=6.0.2",
     "requests>=2.25.0",

tests_integ/cli/runtime/test_simple_agent.py

@@ -1,7 +1,9 @@
+import json
 import logging
 import textwrap
 from typing import List
 
+import boto3
 from click.testing import Result
 
 from tests_integ.cli.runtime.base_test import BaseCLIRuntimeTest, CommandInvocation
@@ -17,6 +19,12 @@ class TestSimpleAgent(BaseCLIRuntimeTest):
     """
 
     def setup(self):
+        # Extract role name from ARN if provided
+        if TEST_ROLE:
+            self.role_name = TEST_ROLE.split("/")[-1]
+        else:
+            self.role_name = None
+
         self.agent_file = "agent.py"
         self.requirements_file = "requirements.txt"
 
@@ -43,6 +51,64 @@ async def agent_invocation(payload):
             """).strip()
             file.write(content)
 
+    def _setup_role_trust_policy(self):
+        """
+        Ensure the IAM role has the required trust relationship with Bedrock.
+        """
+        try:
+            iam_client = boto3.client("iam")
+
+            # Get current trust policy
+            response = iam_client.get_role(RoleName=self.role_name)
+            current_policy = response["Role"]["AssumeRolePolicyDocument"]
+
+            # Check if bedrock is already a trusted service
+            bedrock_trusted = False
+            for statement in current_policy.get("Statement", []):
+                principal = statement.get("Principal", {})
+                service = principal.get("Service", [])
+                if isinstance(service, str):
+                    service = [service]
+                if "bedrock.amazonaws.com" in service:
+                    bedrock_trusted = True
+                    break
+
+            # Add bedrock trust if needed
+            if not bedrock_trusted:
+                logger.info("Adding bedrock.amazonaws.com to trust policy for role %s", self.role_name)
+
+                # Copy the existing policy and add bedrock
+                if len(current_policy.get("Statement", [])) > 0:
+                    # Add to existing policy
+                    new_statement = {
+                        "Effect": "Allow",
+                        "Principal": {"Service": "bedrock.amazonaws.com"},
+                        "Action": "sts:AssumeRole",
+                    }
+                    current_policy["Statement"].append(new_statement)
+                else:
+                    # Create new policy
+                    current_policy = {
+                        "Version": "2012-10-17",
+                        "Statement": [
+                            {
+                                "Effect": "Allow",
+                                "Principal": {"Service": "bedrock.amazonaws.com"},
+                                "Action": "sts:AssumeRole",
+                            }
+                        ],
+                    }
+
+                # Update the role
+                iam_client.update_assume_role_policy(RoleName=self.role_name, PolicyDocument=json.dumps(current_policy))
+                logger.info("Updated trust policy for role %s", self.role_name)
+            else:
+                logger.info("Role %s already trusts bedrock.amazonaws.com", self.role_name)
+
+        except Exception as e:
+            logger.error("Error updating role trust policy: %s", str(e))
+            raise
+
     def get_command_invocations(self) -> List[CommandInvocation]:
         configure_invocation = CommandInvocation(
             command=[
@@ -85,7 +151,13 @@ def validate_configure(self, result: Result):
 
         assert "Configuration Success" in output
         assert "Agent Name: agent" in output
-        assert TEST_ROLE in output
+
+        # Handle both explicit role and auto-create
+        if TEST_ROLE:
+            assert TEST_ROLE in output
+        else:
+            assert "Auto-create" in output or "Execution Role:" in output
+
         assert "Authorization: IAM" in output
         assert ".bedrock_agentcore.yaml" in output
 

tests_integ/utils/config.py

@@ -1,4 +1,4 @@
 import os
 
-TEST_ROLE = os.getenv("AGENTCORE_TEST_ROLE", default="Admin")
+TEST_ROLE = None
 TEST_ECR = os.getenv("AGENTCORE_TEST_ECR", default="auto")