Mock secret manager with pytest

Question

I'm using the default Lambda function to rotate our Aurora password in AWS Code here: https://github.com/aws-samples/aws-secrets-manager-rotation-lambdas/blob/master/SecretsManagerRDSMariaDBRotationSingleUser/lambda_function.py

I have to test this code before it's deployed however I'm not sure how to do it. Can anybody help? I know the code is probably completely wrong, but just need some guidance. I want to test the following function with Pytest.

def test_secret(service_client, arn, token):
    """Args:
     service_client (client): The secrets manager service client
     arn (string): The secret ARN or other identifier
     token (string): The ClientRequestToken associated with the secret version
   Raises:
     ResourceNotFoundException: If the secret with the specified arn and stage does not exist
     ValueError: If the secret is not valid JSON or valid credentials are found to login to the database
     KeyError: If the secret json does not contain the expected keys
 """
    # Try to login with the pending secret, if it succeeds, return
    conn = get_connection(get_secret_dict(service_client, arn, "AWSPENDING", token))
    if conn:
        # This is where the lambda will validate the user's permissions. Uncomment/modify the below lines to
        # tailor these validations to your needs
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT NOW()")
                conn.commit()
        finally:
            conn.close()

        logger.info("testSecret: Successfully signed into MariaDB DB with AWSPENDING secret in %s." % arn)
        return
    else:
        logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn)
        raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn)

import lambda_function.py as testpass
import boto3
import moto import mock_secretsmanager
#Not sure where to get these values from to mock"
token = "akd93939-383838-999388"
arn = "secret-arn"
token = "9393939302931883487"

@mock_secretsmanager
def test_testsecret(mock_secret_manager):
    conn = boto3.client("secretsmanager", region_name="us-east-1")

    test = testpass.test_secret("secretsmanager", arn, token)
    assert test

user_9926721 · Answer

You can mock nested functions using mocking functionality. I renamed the function from test_secret to secret_test because pytest is not happy with that name:

import uuid
from unittest.mock import patch, call, MagicMock

import boto3
import pytest
from moto import mock_secretsmanager

from lambda_function import secret_test

class TestSecret:
    TEST_SECRET_DICT = {'engine': 'mariadb', 'username': 'user123', 'password': 'test_pass', 'host': 'localhost'}
    TEST_ARN = "secret-arn"

    @pytest.fixture
    def mock_get_secret_dict(self):
        with patch('lambda_function.get_secret_dict') as mock:
            mock.return_value = self.TEST_SECRET_DICT
            yield mock

    @pytest.fixture
    def mock_get_connection(self):
        with patch('lambda_function.get_connection') as mock:
            yield mock

    @mock_secretsmanager
    def test_secret_test(self, mock_get_secret_dict, mock_get_connection):
        mock_cursor = MagicMock()
        mock_get_connection.return_value.cursor.return_value.__enter__.return_value = mock_cursor

        request_token = str(uuid.uuid4())

        sm_client = boto3.client("secretsmanager", region_name="us-east-1")

        result = secret_test(sm_client, self.TEST_ARN, request_token)

        assert result is None
        assert mock_get_secret_dict.call_args == call(
            sm_client, self.TEST_ARN, 'AWSPENDING', request_token
        )
        assert mock_get_connection.call_args == call(self.TEST_SECRET_DICT)
        assert mock_get_connection.return_value.method_calls == [
            call.cursor(),
            call.commit(),
            call.close()
        ]
        assert mock_cursor.method_calls == [call.execute('SELECT NOW()')]

Mock secret manager with pytest

1 Answers