standardize parameter names across Vault methods

Author: debussyman

tests/test_vault.py

@@ -350,7 +350,7 @@ def test_encrypt_success(
         plaintext = "Hello, World!"
         context = KeyContext({"key": "test-key"})
 
-        encrypted_data = self.vault.encrypt(data=plaintext, context=context)
+        encrypted_data = self.vault.encrypt(data=plaintext, key_context=context)
 
         # Verify create_data_key was called
         assert request_kwargs["method"] == "post"
@@ -372,7 +372,7 @@ def test_encrypt_with_associated_data(
         associated_data = "additional-context"
 
         encrypted_data = self.vault.encrypt(
-            data=plaintext, context=context, associated_data=associated_data
+            data=plaintext, key_context=context, associated_data=associated_data
         )
 
         # Verify we got encrypted data back
@@ -393,7 +393,7 @@ def test_decrypt_success(self, mock_data_key, capture_and_mock_http_client_reque
 
         plaintext = "Hello, World!"
         context = KeyContext({"key": "test-key"})
-        encrypted_data = self.vault.encrypt(data=plaintext, context=context)
+        encrypted_data = self.vault.encrypt(data=plaintext, key_context=context)
 
         # Now mock decrypt_data_key for decryption
         capture_and_mock_http_client_request(self.http_client, mock_data_key, 200)
@@ -422,7 +422,7 @@ def test_decrypt_with_associated_data(
         context = KeyContext({"key": "test-key"})
         associated_data = "additional-context"
         encrypted_data = self.vault.encrypt(
-            data=plaintext, context=context, associated_data=associated_data
+            data=plaintext, key_context=context, associated_data=associated_data
         )
 
         # Now mock decrypt_data_key for decryption
@@ -448,7 +448,7 @@ def test_encrypt_decrypt_roundtrip(
         context = KeyContext({"env": "test", "service": "vault"})
 
         # Encrypt the data
-        encrypted_data = self.vault.encrypt(data=plaintext, context=context)
+        encrypted_data = self.vault.encrypt(data=plaintext, key_context=context)
 
         # Mock decrypt_data_key for decryption
         capture_and_mock_http_client_request(self.http_client, mock_data_key, 200)

workos/vault.py

@@ -28,7 +28,7 @@
 class VaultModule(Protocol):
     def read_object(self, *, object_id: str) -> VaultObject:
         """
-        Get a Vault object with the decrypted value.
+        Get a Vault object with the value decrypted.
 
         Kwargs:
             object_id (str): The unique identifier for the object.
@@ -81,12 +81,12 @@ def create_object(
         key_context: KeyContext,
     ) -> ObjectMetadata:
         """
-        Create a new Vault object.
+        Create a new Vault encrypted object.
 
         Kwargs:
             name (str): The name of the object.
             value (str): The value to encrypt and store.
-            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
 
         Returns:
             VaultObject: The created vault object.
@@ -119,7 +119,7 @@ def delete_object(
         object_id: str,
     ) -> None:
         """
-        Permanently delete a Vault encrypted object.
+        Permanently delete a Vault encrypted object. Warning: this cannont be undone.
 
         Kwargs:
             object_id (str): The unique identifier for the object.
@@ -132,7 +132,7 @@ def create_data_key(self, *, key_context: KeyContext) -> DataKeyPair:
         The encrypted data key MUST be stored by the application, as it cannot be retrieved after generation.
 
         Kwargs:
-            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
         """
         ...
 
@@ -157,7 +157,11 @@ def decrypt_data_key(
         ...
 
     def encrypt(
-        self, *, data: str, context: KeyContext, associated_data: Optional[str] = None
+        self,
+        *,
+        data: str,
+        key_context: KeyContext,
+        associated_data: Optional[str] = None,
     ) -> str:
         """
         Encrypt data locally using AES-GCM with a data key derived from the provided context.
@@ -168,7 +172,7 @@ def encrypt(
 
         Kwargs:
             data (str): The plaintext data to encrypt.
-            context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use for key derivation.
+            key_context (KeyContext): A set of key-value dictionary pairs that determines which root keys to use when encrypting data.
             associated_data (str): Additional authenticated data (AAD) that will be authenticated but not encrypted. (Optional)
 
         Returns:
@@ -382,9 +386,13 @@ def decrypt_data_key(
         )
 
     def encrypt(
-        self, *, data: str, context: KeyContext, associated_data: Optional[str] = None
+        self,
+        *,
+        data: str,
+        key_context: KeyContext,
+        associated_data: Optional[str] = None,
     ) -> str:
-        key_pair = self.create_data_key(key_context=context)
+        key_pair = self.create_data_key(key_context=key_context)
 
         key = self._base64_to_bytes(key_pair.data_key.key)
         key_blob = self._base64_to_bytes(key_pair.encrypted_keys)