refactoring and fixes

Author: RitvikKapila

examples/src/multithreading/__init__.py

@@ -1,11 +1,12 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """init file for multi-threading examples."""
+import time
+
 from aws_cryptographic_materialproviders.mpl.references import IKeyring
 from typing import Dict  # noqa pylint: disable=wrong-import-order
 
 import aws_encryption_sdk
-import time
 
 
 def encrypt_and_decrypt_with_keyring(

examples/test/multithreading/test_i_raw_aes_keyring_multithreaded_example.py

@@ -11,16 +11,11 @@
 from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds
 from ...src.multithreading.raw_aes_keyring import create_keyring
 
-import time
-
 pytestmark = [pytest.mark.examples]
 
 
-def test_encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=16, duration=60):
+def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60):
     """Helper function for multi-threaded encrypt and decrypt using a keyring for fixed n_threads and duration."""
-    print(n_threads, duration)
-    start_time = time.time()
-    print('start_time', start_time)
     keyring = create_keyring()
     plaintext_data = b"Hello World"
     client = aws_encryption_sdk.EncryptionSDKClient(
@@ -36,15 +31,10 @@ def test_encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=16, dur
 
         for future in as_completed(thread_futures):
             future.result()
-    end_time = time.time()
-    print('end_time', end_time)
-    print('duration', end_time - start_time)
-
-
-# def test_encrypt_and_decrypt_with_keyring_multithreaded(n_threads_list: list = [4, 16, 64], duration_list: list = [2, 10, 60]):
-#     """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
-#     print('hello', n_threads_list, duration_list)
-#     for n in n_threads_list:
-#         for d in duration_list:
-#             print(n, d, time.time())
-#             encrypt_and_decrypt_with_keyring_helper(n_threads=n, duration=d)
+
+
+def test_encrypt_and_decrypt_with_keyring_multithreaded(n_threads_list: list = [1, 4, 16, 64], duration_list: list = [2, 10, 60]):
+    """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
+    for n in n_threads_list:
+        for d in duration_list:
+            encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d)

examples/test/multithreading/test_i_raw_rsa_keyring_multithreaded_example.py

@@ -8,14 +8,14 @@
 import aws_encryption_sdk
 from aws_encryption_sdk import CommitmentPolicy
 
-from ...src.multithreading import encrypt_and_decrypt_with_keyring
+from ...src.multithreading import run_encrypt_and_decrypt_with_keyring_for_duration_seconds
 from ...src.multithreading.raw_rsa_keyring import create_keyring, generate_rsa_keys
 
 pytestmark = [pytest.mark.examples]
 
 
-def test_encrypt_and_decrypt_with_keyring(n_threads=10):
-    """Test function for multi-threaded encrypt and decrypt using the Raw RSA Keyring example."""
+def encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=64, duration=60):
+    """Helper function for multi-threaded encrypt and decrypt using a keyring for fixed n_threads and duration."""
     public_key, private_key = generate_rsa_keys()
     keyring = create_keyring(public_key=public_key, private_key=private_key)
     plaintext_data = b"Hello World"
@@ -24,12 +24,18 @@ def test_encrypt_and_decrypt_with_keyring(n_threads=10):
     )
 
     with ThreadPoolExecutor(max_workers=n_threads) as executor:
-        thread_futures = {executor.submit(encrypt_and_decrypt_with_keyring,
+        thread_futures = {executor.submit(run_encrypt_and_decrypt_with_keyring_for_duration_seconds,
                                           plaintext_data=plaintext_data,
                                           keyring=keyring,
-                                          client=client): i for i in range(n_threads)}
+                                          client=client,
+                                          duration=duration): i for i in range(n_threads)}
 
         for future in as_completed(thread_futures):
-            decrypted_plaintext_data = future.result()
-            assert decrypted_plaintext_data == plaintext_data, \
-                "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
+            future.result()
+
+
+def test_encrypt_and_decrypt_with_keyring_multithreaded(n_threads_list: list = [1, 4, 16, 64], duration_list: list = [2, 10, 60]):
+    """Test function for multi-threaded encrypt and decrypt using a keyring for different n_threads and duration."""
+    for n in n_threads_list:
+        for d in duration_list:
+            encrypt_and_decrypt_with_keyring_multithreaded_helper(n_threads=n, duration=d)