chore(aws_kms): AWS KMS keyring / key-provider test

Author: RitvikKapila

performance_tests/consolidate_results.py

@@ -15,12 +15,14 @@ def calculate_statistics(_csv_file):
 
     # Calculate statistics
     if data:
+        data = np.sort(data)
         _total_entries = len(data)
         _average = sum(data) / _total_entries
+        _trimmed_average_99 = np.mean(data[0:int(0.99 * len(data))])
         _minimum = min(data)
         _maximum = max(data)
-        _perc_99 = np.percentile(data, 99) 
-        return _total_entries, _average, _minimum, _maximum, _perc_99
+        _perc_99 = np.percentile(data, 99)
+        return _total_entries, _average, _trimmed_average_99, _minimum, _maximum, _perc_99
 
     return None
 
@@ -33,10 +35,11 @@ def calculate_statistics(_csv_file):
 
     statistics = calculate_statistics(args.csv_file)
     if statistics:
-        total_entries, average, minimum, maximum, perc_99 = statistics
+        total_entries, average, trimmend_average_99, minimum, maximum, perc_99 = statistics
         print("CSV File:", args.csv_file)
         print("Total Entries:", total_entries)
         print("Average:", average)
+        print("99th percentile trimmed average:", trimmend_average_99)
         print("Minimum:", minimum)
         print("Maximum:", maximum)
         print("99th percentile:", perc_99)

performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/aws_kms_keyring.py

@@ -42,6 +42,49 @@ def create_keyring(
     return keyring
 
 
+def create_kms_client():
+    """Create an AWS KMS client.
+
+    Usage: create_kms_client()
+    """
+    # Create a boto3 client for KMS.
+    kms_client = boto3.client('kms', region_name="us-west-2")
+
+    return kms_client
+
+
+def create_keyring_given_kms_client(
+    kms_key_id: str,
+    kms_client: boto3.client,
+):
+    """Demonstrate how to create an AWS KMS keyring with given KMS client.
+
+    Usage: create_keyring(kms_key_id, kms_client)
+    :param kms_key_id: KMS Key identifier for the KMS key you want to use.
+    :type kms_key_id: string
+    :param kms_client: boto3 client for KMS.
+    :type kms_client: boto3.client
+
+    For more information on KMS Key identifiers, see
+    https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
+    """
+    # Create a KMS keyring
+    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+        config=MaterialProvidersConfig()
+    )
+
+    keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
+        kms_key_id=kms_key_id,
+        kms_client=kms_client
+    )
+
+    keyring: IKeyring = mat_prov.create_aws_kms_keyring(
+        input=keyring_input
+    )
+
+    return keyring
+
+
 def encrypt_using_keyring(
     plaintext_data: bytes,
     keyring: IKeyring

performance_tests/test/keyrings/test_aws_kms_keyring.py

@@ -9,6 +9,8 @@
 
 from aws_encryption_sdk_performance_tests.keyrings.aws_kms_keyring import (
     create_keyring,
+    create_kms_client,
+    create_keyring_given_kms_client,
     decrypt_using_keyring,
     encrypt_using_keyring,
 )
@@ -46,6 +48,38 @@ def create(
     PerfTestUtils.print_time_list_to_csv(time_list, output_file)
 
 
+@click.group()
+def create_kms_keyring_given_kms_client():
+    """Click group helper function"""
+
+
+@create_kms_keyring_given_kms_client.command()
+@click.option('--kms_key_id',
+              default='arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f')
+@click.option('--n_iters',
+              default=PerfTestUtils.DEFAULT_N_ITERS)
+@click.option('--output_file',
+              default='kms_keyring_create_given_kms_client')
+def create_given_kms_client(
+    kms_key_id: str,
+    n_iters: int,
+    output_file: str
+):
+    """Performance test for the create_keyring function."""
+    kms_client = create_kms_client()
+    time_list = []
+    for _ in tqdm(range(n_iters)):
+        curr_time = time.time()
+
+        create_keyring_given_kms_client(kms_key_id, kms_client)
+
+        # calculate elapsed time in milliseconds
+        elapsed_time = (time.time() - curr_time) * 1000
+        time_list.append(elapsed_time)
+
+    PerfTestUtils.print_time_list_to_csv(time_list, output_file)
+
+
 @click.group()
 def encrypt_kms_keyring():
     """Click group helper function"""
@@ -125,6 +159,7 @@ def decrypt(
 
 
 kms_keyring_test = click.CommandCollection(sources=[create_kms_keyring,
+                                                    create_kms_keyring_given_kms_client,
                                                     encrypt_kms_keyring,
                                                     decrypt_kms_keyring])