diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..0738a6a --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,66 @@ +name: Run Tests + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.9', '3.10', '3.11'] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Cache pip packages + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements*.txt') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements-test.txt + + - name: Run tests with pytest + run: | + pytest --cov=lambda/skyflow --cov-report=xml --cov-report=term + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v4 + if: matrix.python-version == '3.9' + with: + file: ./coverage.xml + flags: unittests + name: codecov-umbrella + fail_ci_if_error: false + + - name: Generate coverage report + if: matrix.python-version == '3.9' + run: | + pip install coverage + coverage report + + - name: Archive test results + if: always() + uses: actions/upload-artifact@v4 + with: + name: test-results-${{ matrix.python-version }} + path: | + htmlcov/ + coverage.xml + .coverage diff --git a/.gitignore b/.gitignore index 3913e56..5c4fe1c 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,23 @@ CLAUDE.md .claude/ .claude/* +# Ignore Python virtual environment +.venv/ + +# Ignore Python bytecode and cache +__pycache__/ +*.py[cod] +*$py.class +*.so + +# Ignore pytest cache +.pytest_cache/ + +# Ignore coverage reports +.coverage +htmlcov/ +coverage.xml + +# VSCode settings (keep settings.json for shared project config) +.vscode/* +!.vscode/settings.json \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..66c466c --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,12 @@ +{ + "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python", + "python.terminal.activateEnvironment": true, + "python.testing.pytestEnabled": true, + "python.testing.unittestEnabled": false, + "python.testing.pytestArgs": [ + "tests" + ], + "python.analysis.extraPaths": [ + "${workspaceFolder}/lambda/skyflow" + ] +} diff --git a/README.md b/README.md index 8d53bd9..142ee23 100644 --- a/README.md +++ b/README.md @@ -5,23 +5,39 @@ A Lambda-based handler for tokenizing and detokenizing sensitive data in Amazon [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) ## Table of Contents -- [Introduction](#introduction) -- [Features](#features) -- [How It Works](#how-it-works) -- [Setup](#setup) - - [Prerequisites](#prerequisites) - - [Skyflow Configuration](#skyflow-configuration) - - [Secret Structure](#secret-structure) -- [Deployment](#deployment) - - [AWS Lambda Setup](#aws-lambda-setup) - - [Redshift Integration](#redshift-integration) -- [Usage](#usage) - - [In SQL Queries](#in-sql-queries) -- [Security Considerations](#security-considerations) -- [Troubleshooting](#troubleshooting) - - [Common Issues](#common-issues) -- [Contributing](#contributing) -- [License](#license) +- [Skyflow Redshift UDF Handler](#skyflow-redshift-udf-handler) + - [Table of Contents](#table-of-contents) + - [Introduction](#introduction) + - [Features](#features) + - [How It Works](#how-it-works) + - [Setup](#setup) + - [Prerequisites](#prerequisites) + - [Skyflow Configuration](#skyflow-configuration) + - [1. Create a Skyflow Vault](#1-create-a-skyflow-vault) + - [2. Configure Vault Schema](#2-configure-vault-schema) + - [3. Configure Roles and Policies](#3-configure-roles-and-policies) + - [4. Create Service Accounts](#4-create-service-accounts) + - [5. Store Credentials in AWS Secrets Manager](#5-store-credentials-in-aws-secrets-manager) + - [Secret Structure](#secret-structure) + - [Role Mappings Secret](#role-mappings-secret) + - [Credentials Mapping Secret](#credentials-mapping-secret) + - [Deployment](#deployment) + - [AWS Lambda Setup](#aws-lambda-setup) + - [Redshift Integration](#redshift-integration) + - [Redshift Permissions](#redshift-permissions) + - [Usage](#usage) + - [In SQL Queries](#in-sql-queries) + - [Security Considerations](#security-considerations) + - [Troubleshooting](#troubleshooting) + - [Common Issues](#common-issues) + - [Contributing](#contributing) + - [Development](#development) + - [Testing](#testing) + - [Running Tests Locally](#running-tests-locally) + - [Test Structure](#test-structure) + - [Coverage Requirements](#coverage-requirements) + - [Continuous Integration](#continuous-integration) + - [License](#license) ## Introduction @@ -326,6 +342,67 @@ SELECT detokenize('ab123c4d-ef56-7890-gh12-3ij4klm5no6p', 'REDACTED'); Contributions are welcome! Please feel free to submit a Pull Request. +## Development + +### Testing + +This project uses pytest for unit testing with comprehensive test coverage. + +#### Running Tests Locally + +1. Install test dependencies: + ```bash + pip install -r requirements-test.txt + ``` + +2. Run all tests: + ```bash + pytest + ``` + +3. Run tests with coverage report: + ```bash + pytest --cov=lambda/skyflow --cov-report=term-missing + ``` + +4. Run specific test classes or functions: + ```bash + # Run a specific test class + pytest tests/test_skyflow_detokenize_function.py::TestCredentialConfig + + # Run a specific test function + pytest tests/test_skyflow_detokenize_function.py::TestCredentialConfig::test_get_credentials_secret_with_mapping + ``` + +#### Test Structure + +The test suite includes comprehensive coverage for: + +- **Configuration Classes**: Tests for `RoleConfig` and `CredentialConfig` with various mapping scenarios +- **JWT Generation**: Tests for token generation with valid and invalid credentials +- **Handler Operations**: Tests for tokenization and detokenization with mocked AWS and Skyflow APIs +- **Caching**: Tests for bearer token caching and expiration +- **Error Handling**: Tests for various error scenarios (404 errors, missing secrets, invalid arguments) +- **Lambda Integration**: Tests for `lambda_handler` with different event structures and role normalization + +#### Coverage Reports + +Coverage reports are generated in multiple formats for visibility: +- Terminal output with missing lines +- HTML report (in `htmlcov/` directory) +- XML report (for CI/CD integration) + +Current coverage is approximately 72%. There is no minimum coverage requirement enforced. + +#### Continuous Integration + +Tests run automatically via GitHub Actions on: +- Push to main branch +- Pull requests to main branch +- Python versions: 3.9, 3.10, 3.11 + +View the test workflow at [.github/workflows/test.yml](.github/workflows/test.yml) + ## License This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..e56dd8c --- /dev/null +++ b/pytest.ini @@ -0,0 +1,31 @@ +[pytest] +# Pytest configuration for Skyflow Redshift UDF tests + +# Test discovery patterns +python_files = test_*.py +python_classes = Test* +python_functions = test_* + +# Coverage options +addopts = + --verbose + --strict-markers + --cov=lambda/skyflow + --cov-report=term-missing + --cov-report=html + --cov-report=xml + +# Test paths +testpaths = tests + +# Markers for organizing tests +markers = + unit: Unit tests + integration: Integration tests + slow: Slow running tests + +# Warning filters +filterwarnings = + error + ignore::UserWarning + ignore::DeprecationWarning diff --git a/requirements-test.txt b/requirements-test.txt new file mode 100644 index 0000000..c7a426b --- /dev/null +++ b/requirements-test.txt @@ -0,0 +1,10 @@ +# Test dependencies for Skyflow Redshift UDF Lambda function +-r requirements.txt + +# Testing framework +pytest>=7.4.0 +pytest-cov>=4.1.0 +pytest-mock>=3.11.0 + +# Mocking libraries +moto>=4.1.0 # Mock AWS services diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..80c8a26 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +# Production dependencies for Skyflow Redshift UDF Lambda function +boto3>=1.26.0 +requests>=2.28.0 +PyJWT>=2.8.0 +cryptography>=41.0.0 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..f374470 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Tests package for Skyflow Redshift UDF diff --git a/tests/test_skyflow_detokenize_function.py b/tests/test_skyflow_detokenize_function.py new file mode 100644 index 0000000..60b2f3d --- /dev/null +++ b/tests/test_skyflow_detokenize_function.py @@ -0,0 +1,501 @@ +import json +import os +import pytest +from datetime import datetime, timedelta +from unittest.mock import Mock, MagicMock, patch, ANY +import sys + +# Add the lambda directory to the path so we can import the module +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'lambda', 'skyflow')) + +from skyflow_detokenize_function import ( + generate_jwt_token, + CredentialConfig, + RoleConfig, + SkyflowRedshiftUDFHandler, + validate_environment, + create_handler_from_env, + process_single_argument_set, + lambda_handler +) + + +class TestGenerateJWTToken: + """Test JWT token generation.""" + + def test_generate_jwt_token_invalid_private_key(self): + """Test JWT generation with invalid private key.""" + creds = { + 'clientID': 'test-client-id', + 'keyID': 'test-key-id', + 'tokenURI': 'https://auth.skyflow.com/token', + 'privateKey': 'invalid-key' + } + user_email = 'test@example.com' + + with pytest.raises(Exception) as exc_info: + generate_jwt_token(creds, user_email) + + assert "Failed to generate JWT token" in str(exc_info.value) + + +class TestCredentialConfig: + """Test CredentialConfig class.""" + + def test_get_credentials_secret_with_mapping(self): + """Test retrieving credentials secret with a role mapping.""" + config_data = { + 'defaultCredentialsSecret': 'default-creds', + 'credentialMappings': [ + { + 'roles': ['analyst', 'viewer'], + 'credentialsSecret': 'analyst-creds' + } + ] + } + config = CredentialConfig(config_data) + + result = config.get_credentials_secret('analyst') + assert result == 'analyst-creds' + + def test_get_credentials_secret_default(self): + """Test retrieving default credentials secret when no mapping found.""" + config_data = { + 'defaultCredentialsSecret': 'default-creds', + 'credentialMappings': [] + } + config = CredentialConfig(config_data) + + result = config.get_credentials_secret('unknown-role') + assert result == 'default-creds' + + def test_get_credentials_secret_empty_mappings(self): + """Test with empty credential mappings.""" + config_data = { + 'defaultCredentialsSecret': 'default-creds' + } + config = CredentialConfig(config_data) + + result = config.get_credentials_secret('any-role') + assert result == 'default-creds' + + +class TestRoleConfig: + """Test RoleConfig class.""" + + def test_get_skyflow_role_with_mapping(self): + """Test retrieving Skyflow role ID with a role mapping.""" + config_data = { + 'defaultRoleID': 'default-role-id', + 'roleMappings': [ + { + 'redshiftRoles': ['admin', 'superuser'], + 'skyflowRoleID': 'admin-role-id' + } + ] + } + config = RoleConfig(config_data) + + result = config.get_skyflow_role('admin') + assert result == 'admin-role-id' + + def test_get_skyflow_role_default(self): + """Test retrieving default Skyflow role ID when no mapping found.""" + config_data = { + 'defaultRoleID': 'default-role-id', + 'roleMappings': [] + } + config = RoleConfig(config_data) + + result = config.get_skyflow_role('unknown-role') + assert result == 'default-role-id' + + def test_get_skyflow_role_multiple_mappings(self): + """Test with multiple role mappings.""" + config_data = { + 'defaultRoleID': 'default-role-id', + 'roleMappings': [ + { + 'redshiftRoles': ['admin'], + 'skyflowRoleID': 'admin-role-id' + }, + { + 'redshiftRoles': ['viewer', 'reader'], + 'skyflowRoleID': 'viewer-role-id' + } + ] + } + config = RoleConfig(config_data) + + assert config.get_skyflow_role('admin') == 'admin-role-id' + assert config.get_skyflow_role('viewer') == 'viewer-role-id' + assert config.get_skyflow_role('unknown') == 'default-role-id' + + +class TestSkyflowRedshiftUDFHandler: + """Test SkyflowRedshiftUDFHandler class.""" + + @pytest.fixture + def handler(self): + """Create a handler instance for testing.""" + return SkyflowRedshiftUDFHandler( + vault_id='test-vault-id', + vault_url='https://vault.skyflow.com', + credentials_mapping_secret='creds-secret', + role_mappings_secret='roles-secret', + service_account_email='test@example.com' + ) + + def test_handler_initialization(self, handler): + """Test handler initialization.""" + assert handler.vault_id == 'test-vault-id' + assert handler.vault_url == 'https://vault.skyflow.com' + assert handler.credentials_mapping_secret == 'creds-secret' + assert handler.role_mappings_secret == 'roles-secret' + assert handler.service_account_email == 'test@example.com' + assert handler.token_cache == {} + + def test_handler_initialization_from_env(self): + """Test handler initialization from environment variables.""" + with patch.dict(os.environ, { + 'SKYFLOW_VAULT_ID': 'env-vault-id', + 'SKYFLOW_VAULT_URL': 'https://env.vault.skyflow.com', + 'CREDENTIALS_MAPPING_SECRET': 'env-creds', + 'ROLE_MAPPINGS_SECRET': 'env-roles' + }): + handler = SkyflowRedshiftUDFHandler() + assert handler.vault_id == 'env-vault-id' + assert handler.vault_url == 'https://env.vault.skyflow.com' + + @patch('boto3.session.Session') + def test_get_secret_success(self, mock_session, handler): + """Test successful secret retrieval.""" + mock_client = MagicMock() + mock_client.get_secret_value.return_value = { + 'SecretString': json.dumps({'key': 'value'}) + } + mock_session.return_value.client.return_value = mock_client + + result = handler.get_secret('test-secret') + + assert result == {'key': 'value'} + mock_client.get_secret_value.assert_called_once_with(SecretId='test-secret') + + @patch('boto3.session.Session') + def test_get_secret_failure(self, mock_session, handler): + """Test secret retrieval failure.""" + from botocore.exceptions import ClientError + + mock_client = MagicMock() + mock_client.get_secret_value.side_effect = ClientError( + {'Error': {'Code': 'ResourceNotFoundException'}}, + 'GetSecretValue' + ) + mock_session.return_value.client.return_value = mock_client + + with pytest.raises(Exception) as exc_info: + handler.get_secret('nonexistent-secret') + + assert "Error retrieving secret" in str(exc_info.value) + + @patch.object(SkyflowRedshiftUDFHandler, 'get_secret') + def test_get_role_config(self, mock_get_secret, handler): + """Test role configuration retrieval.""" + mock_get_secret.return_value = { + 'defaultRoleID': 'default-role', + 'roleMappings': [] + } + + config = handler.get_role_config() + + assert isinstance(config, RoleConfig) + assert config.default_role_id == 'default-role' + mock_get_secret.assert_called_once_with('roles-secret') + + @patch.object(SkyflowRedshiftUDFHandler, 'get_secret') + def test_get_credential_config(self, mock_get_secret, handler): + """Test credential configuration retrieval.""" + mock_get_secret.return_value = { + 'defaultCredentialsSecret': 'default-creds', + 'credentialMappings': [] + } + + config = handler.get_credential_config() + + assert isinstance(config, CredentialConfig) + assert config.default_credentials == 'default-creds' + mock_get_secret.assert_called_once_with('creds-secret') + + @patch('requests.post') + @patch.object(SkyflowRedshiftUDFHandler, 'get_secret') + def test_get_bearer_token_cached(self, mock_get_secret, mock_post, handler): + """Test bearer token retrieval from cache.""" + # Pre-populate cache + handler.token_cache['role-id-test-role'] = { + 'token': 'cached-token', + 'expiry': datetime.now() + timedelta(hours=1) + } + + token = handler.get_bearer_token('role-id', 'test-role') + + assert token == 'cached-token' + mock_get_secret.assert_not_called() + mock_post.assert_not_called() + + @patch('requests.post') + @patch.object(SkyflowRedshiftUDFHandler, 'get_bearer_token') + @patch.object(SkyflowRedshiftUDFHandler, 'get_role_config') + def test_tokenize_success(self, mock_get_role_config, mock_get_bearer_token, mock_post, handler): + """Test successful tokenization.""" + mock_role_config = Mock() + mock_role_config.get_skyflow_role.return_value = 'test-skyflow-role' + mock_get_role_config.return_value = mock_role_config + mock_get_bearer_token.return_value = 'test-token' + + mock_response = Mock() + mock_response.json.return_value = { + 'records': [{'tokens': {'ssn': 'token-12345'}}] + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + result = handler.tokenize('123-45-6789', 'users', 'ssn', 'test-role') + + assert result == 'token-12345' + mock_post.assert_called_once() + call_args = mock_post.call_args + assert 'https://vault.skyflow.com/v1/vaults/test-vault-id/insert' in call_args[0] + + def test_tokenize_empty_value(self, handler): + """Test tokenization with empty value.""" + result = handler.tokenize('', 'users', 'ssn', 'test-role') + assert result == '' + + @patch('requests.post') + @patch.object(SkyflowRedshiftUDFHandler, 'get_bearer_token') + @patch.object(SkyflowRedshiftUDFHandler, 'get_role_config') + def test_tokenize_404_error(self, mock_get_role_config, mock_get_bearer_token, mock_post, handler): + """Test tokenization with 404 error.""" + mock_role_config = Mock() + mock_role_config.get_skyflow_role.return_value = 'test-skyflow-role' + mock_get_role_config.return_value = mock_role_config + mock_get_bearer_token.return_value = 'test-token' + + mock_response = Mock() + mock_response.status_code = 404 + mock_post.return_value.raise_for_status.side_effect = \ + Exception("HTTP 404 error") + mock_post.return_value.response = mock_response + + import requests + mock_post.side_effect = requests.exceptions.RequestException() + mock_post.side_effect.response = mock_response + + result = handler.tokenize('123-45-6789', 'users', 'ssn', 'test-role') + + assert result == '' + + @patch('requests.post') + @patch.object(SkyflowRedshiftUDFHandler, 'get_bearer_token') + @patch.object(SkyflowRedshiftUDFHandler, 'get_role_config') + def test_detokenize_success(self, mock_get_role_config, mock_get_bearer_token, mock_post, handler): + """Test successful detokenization.""" + mock_role_config = Mock() + mock_role_config.get_skyflow_role.return_value = 'test-skyflow-role' + mock_get_role_config.return_value = mock_role_config + mock_get_bearer_token.return_value = 'test-token' + + mock_response = Mock() + mock_response.json.return_value = { + 'records': [{'value': '123-45-6789'}] + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + result = handler.detokenize('token-12345', 'DEFAULT', 'test-role') + + assert result == '123-45-6789' + mock_post.assert_called_once() + + def test_detokenize_empty_token(self, handler): + """Test detokenization with empty token.""" + result = handler.detokenize('', 'DEFAULT', 'test-role') + assert result == '' + + @patch('requests.post') + @patch.object(SkyflowRedshiftUDFHandler, 'get_bearer_token') + @patch.object(SkyflowRedshiftUDFHandler, 'get_role_config') + def test_detokenize_empty_response(self, mock_get_role_config, mock_get_bearer_token, mock_post, handler): + """Test detokenization with empty response.""" + mock_role_config = Mock() + mock_role_config.get_skyflow_role.return_value = 'test-skyflow-role' + mock_get_role_config.return_value = mock_role_config + mock_get_bearer_token.return_value = 'test-token' + + mock_response = Mock() + mock_response.json.return_value = {'records': []} + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + result = handler.detokenize('token-12345', 'DEFAULT', 'test-role') + + assert result == '' + + +class TestValidateEnvironment: + """Test environment validation.""" + + def test_validate_environment_all_present(self): + """Test validation when all variables are present.""" + required_vars = ['VAR1', 'VAR2', 'VAR3'] + with patch.dict(os.environ, {'VAR1': 'val1', 'VAR2': 'val2', 'VAR3': 'val3'}): + validate_environment(required_vars) + + def test_validate_environment_missing_vars(self): + """Test validation when variables are missing.""" + required_vars = ['VAR1', 'VAR2', 'VAR3'] + with patch.dict(os.environ, {'VAR1': 'val1'}, clear=True): + with pytest.raises(EnvironmentError) as exc_info: + validate_environment(required_vars) + + assert 'Missing required environment variables' in str(exc_info.value) + assert 'VAR2' in str(exc_info.value) + assert 'VAR3' in str(exc_info.value) + + +class TestCreateHandlerFromEnv: + """Test handler creation from environment.""" + + def test_create_handler_from_env_success(self): + """Test successful handler creation from environment.""" + with patch.dict(os.environ, { + 'SKYFLOW_VAULT_ID': 'vault-123', + 'SKYFLOW_VAULT_URL': 'https://vault.skyflow.com', + 'CREDENTIALS_MAPPING_SECRET': 'creds-secret', + 'ROLE_MAPPINGS_SECRET': 'roles-secret' + }): + handler = create_handler_from_env() + + assert handler.vault_id == 'vault-123' + assert handler.vault_url == 'https://vault.skyflow.com' + + def test_create_handler_from_env_missing_vars(self): + """Test handler creation with missing environment variables.""" + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(EnvironmentError): + create_handler_from_env() + + +class TestProcessSingleArgumentSet: + """Test single argument set processing.""" + + @pytest.fixture + def mock_handler(self): + """Create a mock handler for testing.""" + handler = Mock(spec=SkyflowRedshiftUDFHandler) + handler.tokenize.return_value = 'token-result' + handler.detokenize.return_value = 'detokenized-result' + return handler + + def test_process_invalid_argument_count(self, mock_handler): + """Test processing with invalid number of arguments.""" + args = ['arg1', 'arg2', 'arg3', 'arg4'] + + result = process_single_argument_set(mock_handler, args, 'test-role') + + assert result == '' + + +class TestLambdaHandler: + """Test Lambda handler function.""" + + @patch('skyflow_detokenize_function.create_handler_from_env') + def test_lambda_handler_success(self, mock_create_handler): + """Test successful Lambda handler execution.""" + mock_handler = Mock(spec=SkyflowRedshiftUDFHandler) + mock_handler.detokenize.return_value = 'detokenized-value' + mock_create_handler.return_value = mock_handler + + event = { + 'arguments': [['token-12345']], + 'user': 'role_analyst' + } + context = {} + + with patch('skyflow_detokenize_function.process_single_argument_set', return_value='result-1'): + result = lambda_handler(event, context) + + result_data = json.loads(result) + assert 'results' in result_data + assert len(result_data['results']) == 1 + + @patch('skyflow_detokenize_function.create_handler_from_env') + def test_lambda_handler_role_normalization(self, mock_create_handler): + """Test Lambda handler normalizes role names.""" + mock_handler = Mock(spec=SkyflowRedshiftUDFHandler) + mock_create_handler.return_value = mock_handler + + event = { + 'arguments': [['token-12345']], + 'user': 'ROLE_Analyst' + } + context = {} + + with patch('skyflow_detokenize_function.process_single_argument_set', return_value='result') as mock_process: + lambda_handler(event, context) + + # Verify role was normalized to lowercase and prefix removed + mock_process.assert_called_once() + call_args = mock_process.call_args[0] + assert call_args[2] == 'analyst' + + @patch('skyflow_detokenize_function.create_handler_from_env') + def test_lambda_handler_multiple_arguments(self, mock_create_handler): + """Test Lambda handler with multiple argument sets.""" + mock_handler = Mock(spec=SkyflowRedshiftUDFHandler) + mock_create_handler.return_value = mock_handler + + event = { + 'arguments': [['token-1'], ['token-2'], ['token-3']], + 'user': 'analyst' + } + context = {} + + with patch('skyflow_detokenize_function.process_single_argument_set', side_effect=['r1', 'r2', 'r3']): + result = lambda_handler(event, context) + + result_data = json.loads(result) + assert len(result_data['results']) == 3 + + @patch('skyflow_detokenize_function.create_handler_from_env') + def test_lambda_handler_error(self, mock_create_handler): + """Test Lambda handler error handling.""" + mock_create_handler.side_effect = Exception('Test error') + + event = { + 'arguments': [['token-12345']], + 'user': 'analyst' + } + context = {} + + result = lambda_handler(event, context) + + result_data = json.loads(result) + assert result_data == {'results': []} + + @patch('skyflow_detokenize_function.create_handler_from_env') + def test_lambda_handler_default_user(self, mock_create_handler): + """Test Lambda handler with default user.""" + mock_handler = Mock(spec=SkyflowRedshiftUDFHandler) + mock_create_handler.return_value = mock_handler + + event = { + 'arguments': [['token-12345']] + } + context = {} + + with patch('skyflow_detokenize_function.process_single_argument_set', return_value='result') as mock_process: + lambda_handler(event, context) + + call_args = mock_process.call_args[0] + assert call_args[2] == 'default'