445 lines
14 KiB
Python
445 lines
14 KiB
Python
"""
|
|
Shared pytest fixtures and helpers for field-level pseudonymisation tests.
|
|
|
|
This module provides:
|
|
- Mock Vault client for testing without real Vault connections
|
|
- Sample data fixtures
|
|
- Configuration fixtures for encryption/decryption operations
|
|
- Helper functions for running ops and managing test Vault storage
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from dagster import build_op_context
|
|
from cryptography.fernet import Fernet
|
|
from hvac.exceptions import InvalidPath, Forbidden
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
from template_code_location.field_level_pseudo_anonymisation.config_models.structured_config import (
|
|
AnonymisePseudonymizeStructuredConfig,
|
|
DepseudonymizeStructuredConfig,
|
|
EncryptConfig,
|
|
DecryptConfig,
|
|
PseudoTechniqueConfig,
|
|
DepseudoTechniqueConfig,
|
|
)
|
|
from template_code_location.field_level_pseudo_anonymisation.ops import (
|
|
anonymize_pseudonymize_structured,
|
|
depseudonymize_structured,
|
|
)
|
|
|
|
|
|
# -------------------------------- Mock Vault Storage ----------------------------------------
|
|
|
|
# In-memory Vault simulation for tests
|
|
_test_vault_storage = {}
|
|
_test_vault_access_control = {} # For simulating access control
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_vault_client():
|
|
"""
|
|
Auto-use fixture that mocks the hvac.Client to avoid real Vault connections.
|
|
Uses an in-memory dict to simulate Vault storage for tests.
|
|
Includes access control simulation for AC3.
|
|
"""
|
|
global _test_vault_storage, _test_vault_access_control
|
|
_test_vault_storage = {} # Reset storage before each test
|
|
_test_vault_access_control = {} # Reset access control
|
|
|
|
def mock_read_secret(path, mount_point):
|
|
"""Mock reading secret from Vault with access control"""
|
|
full_path = f"{mount_point}/{path}"
|
|
|
|
# Check access control first
|
|
if full_path in _test_vault_access_control:
|
|
if not _test_vault_access_control[full_path]:
|
|
raise Forbidden(f"Access denied to secret: {full_path}")
|
|
|
|
if full_path not in _test_vault_storage:
|
|
raise InvalidPath(f"Secret not found: {full_path}")
|
|
return {"data": {"data": {"value": _test_vault_storage[full_path]}}}
|
|
|
|
def mock_create_or_update_secret(path, mount_point, secret):
|
|
"""Mock creating/updating secret in Vault"""
|
|
full_path = f"{mount_point}/{path}"
|
|
_test_vault_storage[full_path] = secret["value"]
|
|
|
|
def mock_delete_metadata(path, mount_point):
|
|
"""Mock deleting secret from Vault"""
|
|
full_path = f"{mount_point}/{path}"
|
|
if full_path in _test_vault_storage:
|
|
del _test_vault_storage[full_path]
|
|
if full_path in _test_vault_access_control:
|
|
del _test_vault_access_control[full_path]
|
|
|
|
with patch("hvac.Client") as mock_client_class:
|
|
mock_instance = MagicMock()
|
|
mock_instance.secrets.kv.v2.read_secret_version.side_effect = mock_read_secret
|
|
mock_instance.secrets.kv.v2.create_or_update_secret.side_effect = (
|
|
mock_create_or_update_secret
|
|
)
|
|
mock_instance.secrets.kv.v2.delete_metadata_and_all_versions.side_effect = (
|
|
mock_delete_metadata
|
|
)
|
|
mock_client_class.return_value = mock_instance
|
|
yield mock_instance
|
|
|
|
|
|
# -------------------------------- Sample Data Fixtures ----------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_df():
|
|
"""
|
|
Fixture providing a sample structured dataset with PII data.
|
|
Represents typical data that requires pseudonymisation and restoration.
|
|
"""
|
|
return pd.DataFrame(
|
|
{
|
|
"id": [1, 2, 3, 4, 5],
|
|
"name": [
|
|
"Alice Smith",
|
|
"Bob Jones",
|
|
"Charlie Brown",
|
|
"David Wilson",
|
|
"Eva Garcia",
|
|
],
|
|
"email": [
|
|
"alice@example.com",
|
|
"bob@example.com",
|
|
"charlie@example.com",
|
|
"david@example.com",
|
|
"eva@example.com",
|
|
],
|
|
"ssn": [
|
|
"123-45-6789",
|
|
"234-56-7890",
|
|
"345-67-8901",
|
|
"456-78-9012",
|
|
"567-89-0123",
|
|
],
|
|
"age": [25, 30, 35, 40, 45],
|
|
"salary": [50000.0, 60000.0, 70000.0, 80000.0, 90000.0],
|
|
"department": ["HR", "IT", "Finance", "IT", "HR"],
|
|
}
|
|
)
|
|
|
|
|
|
# -------------------------------- Configuration Fixtures ----------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def encrypt_config_single_field():
|
|
"""
|
|
Configuration for encrypting a single field (email).
|
|
Used to create pseudonymised data for restoration tests.
|
|
"""
|
|
return AnonymisePseudonymizeStructuredConfig(
|
|
used_function=[
|
|
PseudoTechniqueConfig(
|
|
technique=EncryptConfig(
|
|
type="encrypt",
|
|
columns=["email"],
|
|
key_name="test_restoration_key_single",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def decrypt_config_single_field():
|
|
"""
|
|
Configuration for decrypting a single field (email).
|
|
Used to restore original values.
|
|
"""
|
|
return DepseudonymizeStructuredConfig(
|
|
used_function=[
|
|
DepseudoTechniqueConfig(
|
|
technique=DecryptConfig(
|
|
type="decrypt",
|
|
columns=["email"],
|
|
key_name="test_restoration_key_single",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def encrypt_config_multiple_fields():
|
|
"""
|
|
Configuration for encrypting multiple fields (name, email, ssn).
|
|
Tests restoration of multiple sensitive fields.
|
|
"""
|
|
return AnonymisePseudonymizeStructuredConfig(
|
|
used_function=[
|
|
PseudoTechniqueConfig(
|
|
technique=EncryptConfig(
|
|
type="encrypt",
|
|
columns=["name", "email", "ssn"],
|
|
key_name="test_restoration_key_multi",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def decrypt_config_multiple_fields():
|
|
"""
|
|
Configuration for decrypting multiple fields (name, email, ssn).
|
|
"""
|
|
return DepseudonymizeStructuredConfig(
|
|
used_function=[
|
|
DepseudoTechniqueConfig(
|
|
technique=DecryptConfig(
|
|
type="decrypt",
|
|
columns=["name", "email", "ssn"],
|
|
key_name="test_restoration_key_multi",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def encrypt_config_partial_fields():
|
|
"""
|
|
Configuration for encrypting only some fields (email, ssn).
|
|
Tests partial restoration scenarios.
|
|
"""
|
|
return AnonymisePseudonymizeStructuredConfig(
|
|
used_function=[
|
|
PseudoTechniqueConfig(
|
|
technique=EncryptConfig(
|
|
type="encrypt",
|
|
columns=["email", "ssn"],
|
|
key_name="test_restoration_key_partial",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def decrypt_config_partial_fields():
|
|
"""
|
|
Configuration for decrypting only some fields (email, ssn).
|
|
"""
|
|
return DepseudonymizeStructuredConfig(
|
|
used_function=[
|
|
DepseudoTechniqueConfig(
|
|
technique=DecryptConfig(
|
|
type="decrypt",
|
|
columns=["email", "ssn"],
|
|
key_name="test_restoration_key_partial",
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def authorized_multi_key_scenario():
|
|
"""
|
|
Fixture for testing multi-key authorization scenarios.
|
|
Sets up two keys: one authorized, one denied.
|
|
"""
|
|
clear_vault_key("authorized_key")
|
|
clear_vault_key("unauthorized_key")
|
|
|
|
# Create authorized key by generating it
|
|
authorized_key = Fernet.generate_key().decode()
|
|
set_vault_key("authorized_key", authorized_key)
|
|
|
|
# Create unauthorized key and deny access
|
|
unauthorized_key = Fernet.generate_key().decode()
|
|
set_vault_key("unauthorized_key", unauthorized_key)
|
|
deny_vault_access("unauthorized_key")
|
|
|
|
yield {"authorized": "authorized_key", "unauthorized": "unauthorized_key"}
|
|
|
|
# Cleanup
|
|
clear_vault_key("authorized_key")
|
|
clear_vault_key("unauthorized_key")
|
|
|
|
|
|
@pytest.fixture
|
|
def large_dataset():
|
|
"""
|
|
Fixture providing a large dataset (10,000 rows) for performance testing.
|
|
Reusable across multiple performance tests.
|
|
"""
|
|
return pd.DataFrame(
|
|
{
|
|
"id": range(1, 10001),
|
|
"email": [f"user{i}@example.com" for i in range(1, 10001)],
|
|
"name": [f"User {i}" for i in range(1, 10001)],
|
|
"ssn": [f"{i:03d}-{i:02d}-{i:04d}" for i in range(1, 10001)],
|
|
"age": [20 + (i % 50) for i in range(1, 10001)],
|
|
"salary": [30000.0 + (i * 10) for i in range(1, 10001)],
|
|
"department": [["HR", "IT", "Finance", "Sales"][i % 4] for i in range(1, 10001)],
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def vault_test_keys():
|
|
"""
|
|
Session-scoped fixture to pre-generate test keys for faster test execution.
|
|
Avoids repeated key generation in each test.
|
|
"""
|
|
keys = {f"test_key_{i}": Fernet.generate_key().decode() for i in range(10)}
|
|
|
|
return keys
|
|
|
|
|
|
@pytest.fixture
|
|
def cleanup_test_keys(request):
|
|
"""
|
|
Fixture to automatically cleanup test keys after each test.
|
|
Use with: @pytest.mark.usefixtures("cleanup_test_keys")
|
|
"""
|
|
yield
|
|
|
|
# Cleanup all test keys from mock Vault
|
|
test_keys = [k for k in _test_vault_storage.keys() if "test_" in k]
|
|
for key in test_keys:
|
|
_test_vault_storage.pop(key, None)
|
|
|
|
|
|
# -------------------------------- Helper Functions ----------------------------------------
|
|
|
|
|
|
def config_to_dagster_dict(config):
|
|
"""
|
|
Convert Pydantic config to Dagster-compatible dictionary.
|
|
|
|
For AnonymisePseudonymizeStructuredConfig (uses discriminated Union):
|
|
Pydantic v2 outputs: {'technique': {'type': 'encrypt', 'columns': [...], 'key_name': '...'}}
|
|
Dagster expects: {'technique': {'encrypt': {'columns': [...], 'key_name': '...'}}}
|
|
|
|
For DepseudonymizeStructuredConfig (direct DecryptConfig, no Union):
|
|
Pydantic v2 outputs:
|
|
{'technique': {'type': 'decrypt', 'columns': [...], 'key_name': '...'}}
|
|
Dagster expects: Same flat structure with 'type' field
|
|
|
|
Args:
|
|
config: Pydantic config instance
|
|
(AnonymisePseudonymizeStructuredConfig or
|
|
DepseudonymizeStructuredConfig)
|
|
|
|
Returns:
|
|
dict: Dagster-compatible configuration dictionary
|
|
"""
|
|
from template_code_location.field_level_pseudo_anonymisation.config_models.structured_config import (
|
|
AnonymisePseudonymizeStructuredConfig,
|
|
)
|
|
|
|
config_dict = config.model_dump()
|
|
|
|
# Only convert discriminated unions for AnonymisePseudonymizeStructuredConfig
|
|
# DepseudonymizeStructuredConfig uses direct DecryptConfig (no discriminated union)
|
|
if isinstance(config, AnonymisePseudonymizeStructuredConfig):
|
|
if "used_function" in config_dict:
|
|
for func_config in config_dict["used_function"]:
|
|
if "technique" in func_config:
|
|
technique = func_config["technique"]
|
|
# Pydantic outputs flat dict with 'type' field for discriminated unions
|
|
if isinstance(technique, dict) and "type" in technique:
|
|
# Extract the type discriminator
|
|
technique_type = technique["type"]
|
|
# Create nested structure without the 'type' field
|
|
technique_data = {k: v for k, v in technique.items() if k != "type"}
|
|
# Nest under the discriminator key for Dagster
|
|
func_config["technique"] = {technique_type: technique_data}
|
|
|
|
return config_dict
|
|
|
|
|
|
def run_encrypt_op(config, df):
|
|
"""
|
|
Helper function to execute the anonymize_pseudonymize_structured op.
|
|
|
|
Args:
|
|
config: AnonymisePseudonymizeStructuredConfig instance
|
|
df: Input pandas DataFrame
|
|
|
|
Returns:
|
|
tuple: (result_df, metrics) - Output DataFrame and metrics dict
|
|
"""
|
|
context = build_op_context(op_config=config_to_dagster_dict(config))
|
|
result_df, metrics = anonymize_pseudonymize_structured(context, df=df)
|
|
return result_df.value, metrics.value
|
|
|
|
|
|
def run_decrypt_op(config, df):
|
|
"""
|
|
Helper function to execute the depseudonymize_structured op.
|
|
|
|
Args:
|
|
config: DepseudonymizeStructuredConfig instance
|
|
df: Input pandas DataFrame
|
|
|
|
Returns:
|
|
tuple: (result_df, metrics) - Output DataFrame and metrics dict
|
|
"""
|
|
context = build_op_context(op_config=config_to_dagster_dict(config))
|
|
result_df, metrics = depseudonymize_structured(context, df=df)
|
|
return result_df.value, metrics.value
|
|
|
|
|
|
def clear_vault_key(key_name: str):
|
|
"""
|
|
Helper function to clear a key from the simulated Vault storage for test isolation.
|
|
|
|
Args:
|
|
key_name: Name of the key to delete from Vault
|
|
"""
|
|
full_path = f"secret/PseudonymKeys/{key_name}"
|
|
if full_path in _test_vault_storage:
|
|
del _test_vault_storage[full_path]
|
|
if full_path in _test_vault_access_control:
|
|
del _test_vault_access_control[full_path]
|
|
|
|
|
|
def set_vault_key(key_name: str, key_value: str):
|
|
"""
|
|
Helper function to set a key in the simulated Vault storage.
|
|
|
|
Args:
|
|
key_name: Name of the key
|
|
key_value: Value of the key (Fernet key as string)
|
|
"""
|
|
full_path = f"secret/PseudonymKeys/{key_name}"
|
|
_test_vault_storage[full_path] = key_value
|
|
|
|
|
|
def deny_vault_access(key_name: str):
|
|
"""
|
|
Helper function to deny access to a key for authorization testing (AC3).
|
|
|
|
Args:
|
|
key_name: Name of the key to deny access to
|
|
"""
|
|
full_path = f"secret/PseudonymKeys/{key_name}"
|
|
_test_vault_access_control[full_path] = False
|
|
|
|
|
|
def get_vault_key(key_name: str) -> bytes:
|
|
"""
|
|
Helper function to retrieve a key from the simulated Vault storage.
|
|
|
|
Args:
|
|
key_name: Name of the key to retrieve
|
|
|
|
Returns:
|
|
bytes: The encryption key
|
|
"""
|
|
full_path = f"secret/PseudonymKeys/{key_name}"
|
|
if full_path not in _test_vault_storage:
|
|
raise InvalidPath(f"Key not found: {key_name}")
|
|
return _test_vault_storage[full_path].encode()
|