Fix test_management

Author: aclark4life

tests/encryption_/test_base.py

@@ -1,3 +1,5 @@
+import os
+
 import pymongo
 from bson.binary import Binary
 from django.conf import settings
@@ -19,3 +21,18 @@ def assertEncrypted(self, model, field):
             collection = db[model._meta.db_table]
             data = collection.find_one({}, {field: 1, "_id": 0})
             self.assertIsInstance(data[field], Binary)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        AWS_CREDS = {
+            "accessKeyId": os.environ.get("FLE_AWS_KEY", ""),
+            "secretAccessKey": os.environ.get("FLE_AWS_SECRET", ""),
+        }
+        _USE_AWS_KMS = any(AWS_CREDS.values())
+
+        if _USE_AWS_KMS:
+            self.DEFAULT_KMS_PROVIDER = "aws"
+        else:
+            # Local-only fallback
+            self.DEFAULT_KMS_PROVIDER = "local"

tests/encryption_/test_management.py

@@ -1,9 +1,12 @@
 from io import StringIO
 
 from bson import json_util
+from django.core.exceptions import ImproperlyConfigured
 from django.core.management import call_command
+from django.db import connections
 from django.test import modify_settings
 
+from .models import EncryptionKey
 from .test_base import EncryptionTestCase
 
 
@@ -110,23 +113,23 @@ def test_show_encrypted_fields_map(self):
                 self.assertIn(model_key, command_output)
                 self._compare_output(expected, command_output[model_key])
 
-    # FIXME ValueError: master_key is required for kms_provider: 'aws'
-    #
-    # Get master_key from KMS_CREDENTIALS["aws"]["key"] and pass to create_data_key
-    #
-    # def test_missing_key(self):
-    #     test_key = "encryption__patient.patient_record.ssn"
-    #     msg = (
-    #         f"Encryption key {test_key} not found. Have migrated the "
-    #         "<class 'encryption_.models.PatientRecord'> model?"
-    #     )
-    #     EncryptionKey.objects.filter(key_alt_name=test_key).delete()
-    #     try:
-    #         with self.assertRaisesMessage(ImproperlyConfigured, msg):
-    #             call_command("showencryptedfieldsmap", "--database", "encrypted", verbosity=0)
-    #     finally:
-    #         # Replace the deleted key.
-    #         connections["encrypted"].client_encryption.create_data_key(
-    #             kms_provider="aws",
-    #             key_alt_names=[test_key],
-    #         )
+    def test_missing_key(self):
+        test_key = "encryption__patient.patient_record.ssn"
+        msg = (
+            f"Encryption key {test_key} not found. Have migrated the "
+            "<class 'encryption_.models.PatientRecord'> model?"
+        )
+        EncryptionKey.objects.filter(key_alt_name=test_key).delete()
+        try:
+            with self.assertRaisesMessage(ImproperlyConfigured, msg):
+                call_command("showencryptedfieldsmap", "--database", "encrypted", verbosity=0)
+        finally:
+            # Replace the deleted key.
+            master_key = connections["encrypted"].settings_dict["KMS_CREDENTIALS"][
+                self.DEFAULT_KMS_PROVIDER
+            ]
+            connections["encrypted"].client_encryption.create_data_key(
+                kms_provider=self.DEFAULT_KMS_PROVIDER,
+                master_key=master_key,
+                key_alt_names=[test_key],
+            )