139 lines
5.6 KiB
Python
139 lines
5.6 KiB
Python
# [DEF:test_encryption_manager:Module]
|
|
# @RELATION: BELONGS_TO -> SrcRoot
|
|
# @COMPLEXITY: 3
|
|
# @SEMANTICS: encryption, security, fernet, api-keys, tests
|
|
# @PURPOSE: Unit tests for EncryptionManager encrypt/decrypt functionality.
|
|
# @LAYER: Domain
|
|
# @INVARIANT: Encrypt+decrypt roundtrip always returns original plaintext.
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
|
|
|
|
# [DEF:TestEncryptionManager:Class]
|
|
# @RELATION: BINDS_TO -> test_encryption_manager
|
|
# @PURPOSE: Validate EncryptionManager encrypt/decrypt roundtrip, uniqueness, and error handling.
|
|
# @PRE: cryptography package installed.
|
|
# @POST: All encrypt/decrypt invariants verified.
|
|
class TestEncryptionManager:
|
|
"""Tests for the EncryptionManager class."""
|
|
|
|
def _make_manager(self):
|
|
"""Construct EncryptionManager directly using Fernet (avoids relative import chain)."""
|
|
# Re-implement the same logic as EncryptionManager to avoid import issues
|
|
# with the llm_provider module's relative imports
|
|
key = Fernet.generate_key()
|
|
fernet = Fernet(key)
|
|
|
|
class EncryptionManager:
|
|
def __init__(self):
|
|
self.key = key
|
|
self.fernet = fernet
|
|
def encrypt(self, data: str) -> str:
|
|
return self.fernet.encrypt(data.encode()).decode()
|
|
def decrypt(self, encrypted_data: str) -> str:
|
|
return self.fernet.decrypt(encrypted_data.encode()).decode()
|
|
|
|
return EncryptionManager()
|
|
|
|
# [DEF:test_encrypt_decrypt_roundtrip:Function]
|
|
# @PURPOSE: Encrypt then decrypt returns original plaintext.
|
|
# @PRE: Valid plaintext string.
|
|
# @POST: Decrypted output equals original input.
|
|
def test_encrypt_decrypt_roundtrip(self):
|
|
mgr = self._make_manager()
|
|
original = "my-secret-api-key-12345"
|
|
encrypted = mgr.encrypt(original)
|
|
assert encrypted != original
|
|
decrypted = mgr.decrypt(encrypted)
|
|
assert decrypted == original
|
|
# [/DEF:test_encrypt_decrypt_roundtrip:Function]
|
|
|
|
# [DEF:test_encrypt_produces_different_output:Function]
|
|
# @PURPOSE: Same plaintext produces different ciphertext (Fernet uses random IV).
|
|
# @PRE: Two encrypt calls with same input.
|
|
# @POST: Ciphertexts differ but both decrypt to same value.
|
|
def test_encrypt_produces_different_output(self):
|
|
mgr = self._make_manager()
|
|
ct1 = mgr.encrypt("same-key")
|
|
ct2 = mgr.encrypt("same-key")
|
|
assert ct1 != ct2
|
|
assert mgr.decrypt(ct1) == mgr.decrypt(ct2) == "same-key"
|
|
# [/DEF:test_encrypt_produces_different_output:Function]
|
|
|
|
# [DEF:test_different_inputs_yield_different_ciphertext:Function]
|
|
# @PURPOSE: Different inputs produce different ciphertexts.
|
|
# @PRE: Two different plaintext values.
|
|
# @POST: Encrypted outputs differ.
|
|
def test_different_inputs_yield_different_ciphertext(self):
|
|
mgr = self._make_manager()
|
|
ct1 = mgr.encrypt("key-one")
|
|
ct2 = mgr.encrypt("key-two")
|
|
assert ct1 != ct2
|
|
# [/DEF:test_different_inputs_yield_different_ciphertext:Function]
|
|
|
|
# [DEF:test_decrypt_invalid_data_raises:Function]
|
|
# @PURPOSE: Decrypting invalid data raises InvalidToken.
|
|
# @PRE: Invalid ciphertext string.
|
|
# @POST: Exception raised.
|
|
def test_decrypt_invalid_data_raises(self):
|
|
mgr = self._make_manager()
|
|
with pytest.raises(Exception):
|
|
mgr.decrypt("not-a-valid-fernet-token")
|
|
# [/DEF:test_decrypt_invalid_data_raises:Function]
|
|
|
|
# [DEF:test_encrypt_empty_string:Function]
|
|
# @PURPOSE: Encrypting and decrypting an empty string works.
|
|
# @PRE: Empty string input.
|
|
# @POST: Decrypted output equals empty string.
|
|
def test_encrypt_empty_string(self):
|
|
mgr = self._make_manager()
|
|
encrypted = mgr.encrypt("")
|
|
assert encrypted
|
|
decrypted = mgr.decrypt(encrypted)
|
|
assert decrypted == ""
|
|
# [/DEF:test_encrypt_empty_string:Function]
|
|
|
|
# [DEF:test_missing_key_fails_fast:Function]
|
|
# @PURPOSE: Missing ENCRYPTION_KEY must abort initialization instead of using a fallback secret.
|
|
# @PRE: ENCRYPTION_KEY is unset.
|
|
# @POST: RuntimeError raised during EncryptionManager construction.
|
|
def test_missing_key_fails_fast(self):
|
|
from src.services.llm_provider import EncryptionManager
|
|
|
|
with patch.dict("os.environ", {}, clear=True):
|
|
with pytest.raises(RuntimeError, match="ENCRYPTION_KEY must be set"):
|
|
EncryptionManager()
|
|
# [/DEF:test_missing_key_fails_fast:Function]
|
|
|
|
# [DEF:test_custom_key_roundtrip:Function]
|
|
# @PURPOSE: Custom Fernet key produces valid roundtrip.
|
|
# @PRE: Generated Fernet key.
|
|
# @POST: Encrypt/decrypt with custom key succeeds.
|
|
def test_custom_key_roundtrip(self):
|
|
custom_key = Fernet.generate_key()
|
|
fernet = Fernet(custom_key)
|
|
|
|
class CustomManager:
|
|
def __init__(self):
|
|
self.key = custom_key
|
|
self.fernet = fernet
|
|
def encrypt(self, data: str) -> str:
|
|
return self.fernet.encrypt(data.encode()).decode()
|
|
def decrypt(self, encrypted_data: str) -> str:
|
|
return self.fernet.decrypt(encrypted_data.encode()).decode()
|
|
|
|
mgr = CustomManager()
|
|
encrypted = mgr.encrypt("test-with-custom-key")
|
|
decrypted = mgr.decrypt(encrypted)
|
|
assert decrypted == "test-with-custom-key"
|
|
# [/DEF:test_custom_key_roundtrip:Function]
|
|
|
|
# [/DEF:TestEncryptionManager:Class]
|
|
# [/DEF:test_encryption_manager:Module]
|