d95e4db8f3
* Add AWS Bedrock session token support in API and worker * Add Bedrock auth mode integration tests for AI agents * Split Bedrock integration test env vars for IAM and session creds * cleaning * Add masked Bedrock bearer-token debug info logs * Revert "Add masked Bedrock bearer-token debug info logs" This reverts commit 6b2fc5e7c2d5b1c6db81f416a4439941a084108c. * cleaning
112 lines
4.6 KiB
Python
112 lines
4.6 KiB
Python
"""
|
|
Bedrock auth mode tests for AI agents.
|
|
|
|
Verifies Bedrock works with the following resource auth modes:
|
|
- API key
|
|
- IAM access key + secret key (no session token)
|
|
- Environment credential fallback
|
|
- IAM access key + secret key + session token
|
|
"""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from .conftest import AIAgentTestClient, create_ai_agent_flow
|
|
from .providers import BEDROCK_API_KEY, BEDROCK_ENV, BEDROCK_IAM, BEDROCK_IAM_SESSION
|
|
|
|
|
|
class TestBedrockAuthModes:
|
|
"""Test AI agent Bedrock auth modes."""
|
|
|
|
@pytest.mark.usefixtures("setup_providers")
|
|
def test_bedrock_api_key_resource(self, client: AIAgentTestClient):
|
|
"""Test Bedrock with API key resource auth."""
|
|
if not os.environ.get("BEDROCK_API_KEY"):
|
|
pytest.skip("BEDROCK_API_KEY not set")
|
|
|
|
flow_value = create_ai_agent_flow(
|
|
provider_input_transform=BEDROCK_API_KEY["input_transform"],
|
|
system_prompt="You are a concise assistant.",
|
|
)
|
|
result = client.run_preview_flow(
|
|
flow_value=flow_value,
|
|
args={"user_message": "Reply with one short sentence."},
|
|
)
|
|
|
|
assert isinstance(result.get("output"), str), f"Expected string output, got: {result}"
|
|
assert result.get("output", "").strip(), f"Expected non-empty output, got: {result}"
|
|
|
|
@pytest.mark.usefixtures("setup_providers")
|
|
def test_bedrock_iam_resource_without_session_token(self, client: AIAgentTestClient):
|
|
"""Test Bedrock with IAM access key + secret key auth."""
|
|
if not (
|
|
os.environ.get("BEDROCK_IAM_ACCESS_KEY_ID")
|
|
and os.environ.get("BEDROCK_IAM_SECRET_ACCESS_KEY")
|
|
):
|
|
pytest.skip(
|
|
"BEDROCK_IAM_ACCESS_KEY_ID and BEDROCK_IAM_SECRET_ACCESS_KEY required"
|
|
)
|
|
|
|
flow_value = create_ai_agent_flow(
|
|
provider_input_transform=BEDROCK_IAM["input_transform"],
|
|
system_prompt="You are a concise assistant.",
|
|
)
|
|
result = client.run_preview_flow(
|
|
flow_value=flow_value,
|
|
args={"user_message": "Reply with one short sentence."},
|
|
)
|
|
|
|
assert isinstance(result.get("output"), str), f"Expected string output, got: {result}"
|
|
assert result.get("output", "").strip(), f"Expected non-empty output, got: {result}"
|
|
|
|
@pytest.mark.usefixtures("setup_providers")
|
|
def test_bedrock_environment_resource(self, client: AIAgentTestClient):
|
|
"""Test Bedrock with environment credential fallback."""
|
|
aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
|
|
aws_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
|
|
aws_session_token = os.environ.get("AWS_SESSION_TOKEN")
|
|
|
|
if not (aws_access_key_id and aws_secret_access_key):
|
|
pytest.skip("AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY required for env fallback")
|
|
if aws_access_key_id.startswith("ASIA") and not aws_session_token:
|
|
pytest.skip(
|
|
"AWS_SESSION_TOKEN required for env fallback when AWS_ACCESS_KEY_ID is temporary (ASIA...)"
|
|
)
|
|
|
|
flow_value = create_ai_agent_flow(
|
|
provider_input_transform=BEDROCK_ENV["input_transform"],
|
|
system_prompt="You are a concise assistant.",
|
|
)
|
|
result = client.run_preview_flow(
|
|
flow_value=flow_value,
|
|
args={"user_message": "Reply with one short sentence."},
|
|
)
|
|
|
|
assert isinstance(result.get("output"), str), f"Expected string output, got: {result}"
|
|
assert result.get("output", "").strip(), f"Expected non-empty output, got: {result}"
|
|
|
|
@pytest.mark.usefixtures("setup_providers")
|
|
def test_bedrock_iam_resource_with_session_token(self, client: AIAgentTestClient):
|
|
"""Test Bedrock with IAM access key + secret key + session token auth."""
|
|
if not (
|
|
os.environ.get("BEDROCK_SESSION_ACCESS_KEY_ID")
|
|
and os.environ.get("BEDROCK_SESSION_SECRET_ACCESS_KEY")
|
|
and os.environ.get("BEDROCK_SESSION_TOKEN")
|
|
):
|
|
pytest.skip(
|
|
"BEDROCK_SESSION_ACCESS_KEY_ID, BEDROCK_SESSION_SECRET_ACCESS_KEY, and BEDROCK_SESSION_TOKEN required"
|
|
)
|
|
|
|
flow_value = create_ai_agent_flow(
|
|
provider_input_transform=BEDROCK_IAM_SESSION["input_transform"],
|
|
system_prompt="You are a concise assistant.",
|
|
)
|
|
result = client.run_preview_flow(
|
|
flow_value=flow_value,
|
|
args={"user_message": "Reply with one short sentence."},
|
|
)
|
|
|
|
assert isinstance(result.get("output"), str), f"Expected string output, got: {result}"
|
|
assert result.get("output", "").strip(), f"Expected non-empty output, got: {result}"
|