From 736f8b5a8f5e68c9b6225b99316ac930c6e5cdca Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Mon, 24 Feb 2020 08:24:14 +0000
Subject: [PATCH] Refactor - reuse logic that expects CW log message

---
 tests/test_awslambda/test_lambda.py | 65 ++++++++++-------------------
 1 file changed, 22 insertions(+), 43 deletions(-)

diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py
index 2bd8f4bb3..48539c0e6 100644
--- a/tests/test_awslambda/test_lambda.py
+++ b/tests/test_awslambda/test_lambda.py
@@ -1111,7 +1111,6 @@ def test_create_event_source_mapping():
 @mock_lambda
 @mock_sqs
 def test_invoke_function_from_sqs():
-    logs_conn = boto3.client("logs", region_name="us-east-1")
     sqs = boto3.resource("sqs", region_name="us-east-1")
     queue = sqs.create_queue(QueueName="test-sqs-queue1")
 
@@ -1137,32 +1136,18 @@ def test_invoke_function_from_sqs():
 
     sqs_client = boto3.client("sqs", region_name="us-east-1")
     sqs_client.send_message(QueueUrl=queue.url, MessageBody="test")
-    start = time.time()
-    while (time.time() - start) < 30:
-        result = logs_conn.describe_log_streams(logGroupName="/aws/lambda/testFunction")
-        log_streams = result.get("logStreams")
-        if not log_streams:
-            time.sleep(1)
-            continue
 
-        assert len(log_streams) == 1
-        result = logs_conn.get_log_events(
-            logGroupName="/aws/lambda/testFunction",
-            logStreamName=log_streams[0]["logStreamName"],
-        )
-        for event in result.get("events"):
-            if event["message"] == "get_test_zip_file3 success":
-                return
-        time.sleep(1)
+    expected_msg = "get_test_zip_file3 success"
+    log_group = "/aws/lambda/testFunction"
+    msg_showed_up = wait_for_log_msg(expected_msg, log_group)
 
-    assert False, "Test Failed"
+    assert msg_showed_up, "Message was not found in log_group, so sending an SQS message did not result in a successful Lambda execution"
 
 
 @mock_logs
 @mock_lambda
 @mock_dynamodb2
 def test_invoke_function_from_dynamodb_put():
-    logs_conn = boto3.client("logs", region_name="us-east-1")
     dynamodb = boto3.client("dynamodb", region_name="us-east-1")
     table_name = "table_with_stream"
     table = dynamodb.create_table(
@@ -1197,32 +1182,18 @@ def test_invoke_function_from_dynamodb_put():
     assert response["State"] == "Enabled"
 
     dynamodb.put_item(TableName=table_name, Item={"id": {"S": "item 1"}})
-    start = time.time()
-    while (time.time() - start) < 30:
-        result = logs_conn.describe_log_streams(logGroupName="/aws/lambda/testFunction")
-        log_streams = result.get("logStreams")
-        if not log_streams:
-            time.sleep(1)
-            continue
 
-        assert len(log_streams) == 1
-        result = logs_conn.get_log_events(
-            logGroupName="/aws/lambda/testFunction",
-            logStreamName=log_streams[0]["logStreamName"],
-        )
-        for event in result.get("events"):
-            if event["message"] == "get_test_zip_file3 success":
-                return
-        time.sleep(1)
+    expected_msg = "get_test_zip_file3 success"
+    log_group = "/aws/lambda/testFunction"
+    msg_showed_up = wait_for_log_msg(expected_msg, log_group)
 
-    assert False, "Test Failed"
+    assert msg_showed_up, "Message was not found in log_group, so inserting DynamoDB did not result in a successful Lambda execution"
 
 
 @mock_logs
 @mock_lambda
 @mock_dynamodb2
 def test_invoke_function_from_dynamodb_update():
-    logs_conn = boto3.client("logs", region_name="us-east-1")
     dynamodb = boto3.client("dynamodb", region_name="us-east-1")
     table_name = "table_with_stream"
     table = dynamodb.create_table(
@@ -1263,9 +1234,18 @@ def test_invoke_function_from_dynamodb_update():
         ExpressionAttributeNames={"#attr": "new_attr"},
         ExpressionAttributeValues={":val": {"S": "new_val"}},
     )
+    expected_msg = "get_test_zip_file3 success"
+    log_group = "/aws/lambda/testFunction"
+    msg_showed_up = wait_for_log_msg(expected_msg, log_group)
+
+    assert msg_showed_up, "Message was not found in log_group, so updating DynamoDB did not result in a successful Lambda execution"
+
+
+def wait_for_log_msg(expected_msg, log_group):
+    logs_conn = boto3.client("logs", region_name="us-east-1")
     start = time.time()
     while (time.time() - start) < 30:
-        result = logs_conn.describe_log_streams(logGroupName="/aws/lambda/testFunction")
+        result = logs_conn.describe_log_streams(logGroupName=log_group)
         log_streams = result.get("logStreams")
         if not log_streams:
             time.sleep(1)
@@ -1273,15 +1253,14 @@ def test_invoke_function_from_dynamodb_update():
 
         assert len(log_streams) == 1
         result = logs_conn.get_log_events(
-            logGroupName="/aws/lambda/testFunction",
+            logGroupName=log_group,
             logStreamName=log_streams[0]["logStreamName"],
         )
         for event in result.get("events"):
-            if event["message"] == "get_test_zip_file3 success":
-                return
+            if event["message"] == expected_msg:
+                return True
         time.sleep(1)
-
-    assert False, "Test Failed"
+    return False
 
 
 @mock_logs