diff --git a/bin/run-integ-tests.sh b/bin/run-integ-tests.sh
index f8d5d98..51fb3de 100755
--- a/bin/run-integ-tests.sh
+++ b/bin/run-integ-tests.sh
@@ -18,6 +18,27 @@ status_code=0
 # Configure and start the agent
 ###################################
 
+# Check if IAM user credentials exist
+if [ -z "$AWS_ACCESS_KEY_ID" ] || [ -z "$AWS_SECRET_ACCESS_KEY" ]; then
+    echo "No IAM user credentials found, assuming we are running on CodeBuild pipeline, falling back to IAM role..."
+    
+    # Store the AWS STS assume-role output and extract credentials
+    CREDS=$(aws sts assume-role \
+        --role-arn $Code_Build_Execution_Role_ARN \
+        --role-session-name "session-$(uuidgen)" \
+        --query 'Credentials.[AccessKeyId,SecretAccessKey,SessionToken]' \
+        --output text \
+        --duration-seconds 3600)
+
+    # Parse the output into separate variables
+    read AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN <<< $CREDS
+
+    # Export the variables
+    export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
+else
+    echo "Using provided IAM user credentials..."
+fi
+
 $rootdir/bin/start-agent.sh
 
 ###################################
diff --git a/bin/start-agent.sh b/bin/start-agent.sh
index 0fb09a8..e4b0c7d 100755
--- a/bin/start-agent.sh
+++ b/bin/start-agent.sh
@@ -22,6 +22,7 @@ cd $rootdir/tests/integ/agent
 echo "[AmazonCloudWatchAgent]
 aws_access_key_id = $AWS_ACCESS_KEY_ID
 aws_secret_access_key = $AWS_SECRET_ACCESS_KEY
+aws_session_token = $AWS_SESSION_TOKEN
 " > ./.aws/credentials
 
 echo "[profile AmazonCloudWatchAgent]
diff --git a/tests/canary/agent/container-definitions.json b/tests/canary/agent/container-definitions.json
index 3b21fe9..facf927 100644
--- a/tests/canary/agent/container-definitions.json
+++ b/tests/canary/agent/container-definitions.json
@@ -30,7 +30,7 @@
   },
   {
     "name": "cloudwatch-agent-python",
-    "image": "amazon/cloudwatch-agent:latest",
+    "image": "public.ecr.aws/cloudwatch-agent/cloudwatch-agent:latest",
     "logConfiguration": {
       "logDriver": "awslogs",
       "options": {
@@ -47,4 +47,4 @@
       }
     ]
   }
-]
\ No newline at end of file
+]
diff --git a/tox.ini b/tox.ini
index fe51c57..a8a60fc 100644
--- a/tox.ini
+++ b/tox.ini
@@ -23,6 +23,7 @@ passenv =
     AWS_REGION
     AWS_ACCESS_KEY_ID
     AWS_SECRET_ACCESS_KEY
+    AWS_SESSION_TOKEN
 
 [testenv:flake8]
 basepython = python3.7