From 9b13afb7e6dbda372488af32cd366e0c5abba4d6 Mon Sep 17 00:00:00 2001
From: Wesley Pettit <wppttt@amazon.com>
Date: Thu, 8 Jun 2023 12:28:58 -0700
Subject: [PATCH] load-test: fix failing tests due to role session timeout,
 clean up code

Signed-off-by: Wesley Pettit <wppttt@amazon.com>
---
 buildspec_load_test.yml |  8 --------
 load_tests/load_test.py | 17 +++++++++++++----
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/buildspec_load_test.yml b/buildspec_load_test.yml
index b11022718..3cb45fc2b 100644
--- a/buildspec_load_test.yml
+++ b/buildspec_load_test.yml
@@ -52,12 +52,6 @@ phases:
       - python -m venv venv
       - source venv/bin/activate
       - pip install -r ./load_tests/requirements.txt
-      # AWS configure sts temp credentials
-      - |
-        CREDS=`aws sts assume-role --role-arn ${LOAD_TEST_CFN_ROLE_ARN} --duration-seconds 900 --role-session-name load-test-cfn`
-        aws configure set aws_access_key_id `echo $CREDS | jq -r .Credentials.AccessKeyId`
-        aws configure set aws_secret_access_key `echo $CREDS | jq -r .Credentials.SecretAccessKey`
-        aws configure set aws_session_token `echo $CREDS | jq -r .Credentials.SessionToken`
       - |
         if [ "${PLATFORM}" == "EKS" ]; then
           aws eks update-kubeconfig --name $EKS_CLUSTER_NAME
@@ -65,8 +59,6 @@ phases:
       # Create and set up related testing resources
       - python ./load_tests/load_test.py create_testing_resources
       - source ./load_tests/setup_test_environment.sh
-      # Remove sts temp credentials, delegate responsibility to py script
-      - rm ~/.aws/credentials
       # Run load tests on corresponding platform
       - python ./load_tests/load_test.py ${PLATFORM}
     finally:
diff --git a/load_tests/load_test.py b/load_tests/load_test.py
index 8e1d65802..4c830bbbd 100644
--- a/load_tests/load_test.py
+++ b/load_tests/load_test.py
@@ -200,7 +200,6 @@ def run_ecs_tests():
         session = get_sts_boto_session()
 
         client = session.client('ecs')
-        waiter = client.get_waiter('tasks_stopped')
         
         processes = []
 
@@ -224,7 +223,13 @@ def run_ecs_tests():
         # Validation input type banner
         print(f'\nTest {input_logger["name"]} to {OUTPUT_PLUGIN} in progress...')
 
+    # wait for tasks and validate
+    for input_logger in INPUT_LOGGERS:
         # Wait until task stops and start validation
+        session = get_sts_boto_session()
+
+        client = session.client('ecs')
+        waiter = client.get_waiter('tasks_stopped')
         for throughput in THROUGHPUT_LIST:
             task_arn = names[f'{OUTPUT_PLUGIN}_{throughput}_task_arn']
             waiter.wait(
@@ -266,6 +271,9 @@ def run_ecs_tests():
                 os.environ['LOG_PREFIX'] = resource_resolver.get_destination_s3_prefix(test_configuration["input_configuration"], OUTPUT_PLUGIN)
                 os.environ['DESTINATION'] = 's3'
 
+            # integ test sessions have expired due to long waits for tasks to complete
+            session = get_sts_boto_session()
+
             # Go script environment with sts cred variables
             credentials = session.get_credentials()
             auth_env = {
@@ -286,8 +294,8 @@ def run_ecs_tests():
         for p in processes:
             p["process"].wait()
             stdout, stderr = p["process"].communicate()
-            print(f'raw validator stdout: {stdout}')
-            print(f'raw validator stderr: {stderr}')
+            print(f'{input_logger["name"]} to {OUTPUT_PLUGIN} raw validator stdout: {stdout}')
+            print(f'{input_logger["name"]} to {OUTPUT_PLUGIN} raw validator stderr: {stderr}')
             p["result"] = stdout
         print(f'Test {input_logger["name"]} to {OUTPUT_PLUGIN} complete.')
 
@@ -513,7 +521,8 @@ def get_sts_boto_session():
     # ARN and a role session name.
     assumed_role_object = sts_client.assume_role(
         RoleArn=os.environ["LOAD_TEST_CFN_ROLE_ARN"],
-        RoleSessionName="load-test-cfn"
+        RoleSessionName="load-test-cfn",
+        DurationSeconds=3600
     )
 
     # From the response that contains the assumed role, get the temporary