adding ExecutionTime context manager for benchmarking load tests

Author: malachi-constant

load_tests/_utils.py

@@ -1,5 +1,6 @@
 import random
 from datetime import datetime
+from timeit import default_timer as timer
 from typing import Iterator
 
 import boto3
@@ -11,6 +12,20 @@
 CFN_VALID_STATUS = ["CREATE_COMPLETE", "ROLLBACK_COMPLETE", "UPDATE_COMPLETE", "UPDATE_ROLLBACK_COMPLETE"]
 
 
+class ExecutionTimer:
+    def __init__(self, msg="elapsed time"):
+        self.msg = msg
+
+    def __enter__(self):
+        self.before = timer()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.elapsed_time = round((timer() - self.before), 3)
+        print(f"{self.msg}: {self.elapsed_time:.3f} sec")
+        return None
+
+
 def extract_cloudformation_outputs():
     outputs = {}
     client = boto3.client("cloudformation")

load_tests/test_s3.py

@@ -1,41 +1,39 @@
-import time
-
 import pandas as pd
 import pytest
 
 import awswrangler as wr
 
+from ._utils import ExecutionTimer
+
 
 @pytest.mark.repeat(1)
 @pytest.mark.parametrize("benchmark_time", [150])
 def test_s3_select(benchmark_time):
-    start = time.time()
 
     path = "s3://ursa-labs-taxi-data/2018/1*.parquet"
-    wr.s3.select_query(
-        sql="SELECT * FROM s3object",
-        path=path,
-        input_serialization="Parquet",
-        input_serialization_params={},
-        scan_range_chunk_size=16 * 1024 * 1024,
-    )
-    end = time.time()
+    with ExecutionTimer("elapsed time of wr.s3.select_query()") as timer:
+        wr.s3.select_query(
+            sql="SELECT * FROM s3object",
+            path=path,
+            input_serialization="Parquet",
+            input_serialization_params={},
+            scan_range_chunk_size=16 * 1024 * 1024,
+        )
 
-    elapsed_time = end - start
-    assert elapsed_time < benchmark_time
+    assert timer.elapsed_time < benchmark_time
 
 
-def test_s3_delete_objects(path, path2):
+@pytest.mark.parametrize("benchmark_time", [15])
+def test_s3_delete_objects(path, path2, benchmark_time):
     df = pd.DataFrame({"id": [1, 2, 3]})
     objects_per_bucket = 505
     paths1 = [f"{path}delete-test{i}.json" for i in range(objects_per_bucket)]
     paths2 = [f"{path2}delete-test{i}.json" for i in range(objects_per_bucket)]
     paths = paths1 + paths2
     for path in paths:
         wr.s3.to_json(df, path)
-    start = time.time()
-    wr.s3.delete_objects(path=paths)
-    end = time.time()
-    print(f"elapsed time: {end - start}")
+    with ExecutionTimer("elapsed time of wr.s3.delete_objects()") as timer:
+        wr.s3.delete_objects(path=paths)
+    assert timer.elapsed_time < benchmark_time
     assert len(wr.s3.list_objects(f"{path}delete-test*")) == 0
     assert len(wr.s3.list_objects(f"{path2}delete-test*")) == 0