From 68f762a8b10bf32ddf41d747c5ea8bbee995bc50 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Mon, 29 Apr 2024 12:05:51 -0700 Subject: [PATCH] Move some hooks into tasks --- .../providers/amazon/aws/example_bedrock_knowledge_base.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py b/tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py index f49c3e478eebf..0e8fb0cbef9b0 100644 --- a/tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py +++ b/tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py @@ -149,7 +149,7 @@ def _create_access_policy(name, policy_type, policy): "ResourceType": "index", }, ], - "Principal": [(sts_hook.conn.get_caller_identity()["Arn"]), bedrock_role_arn], + "Principal": [(StsHook().conn.get_caller_identity()["Arn"]), bedrock_role_arn], } ], ) @@ -253,7 +253,7 @@ def copy_data_to_s3(bucket: str): for source in sources: with tempfile.NamedTemporaryFile(mode="w", prefix="") as data_file: urlretrieve(source, data_file.name) - s3_hook.conn.upload_file( + S3Hook().conn.upload_file( Filename=data_file.name, Bucket=bucket, Key=os.path.basename(data_file.name) ) @@ -385,9 +385,6 @@ def delete_opensearch_policies(): aoss_client = OpenSearchServerlessHook().conn bedrock_agent_client = BedrockAgentHook().conn - iam_client = boto3.client("iam") - sts_hook = StsHook() - s3_hook = S3Hook() region_name = boto3.session.Session().region_name