Search code examples
pythonamazon-web-servicesmockingpytestmoto

Python mock AWS SSM


I have written a code that will fetch SSM parameters for me

import boto3
    
client = boto3.client('ssm')
    
def lambda_handler(event, context):
    return client.get_parameter(Name=event["param"], WithDecryption=True)

if __name__ == '__main__':
    print(lambda_handler({"param": "/mypath/password"}, ""))

However, I am not able to write a test case for it I have tried using moto but for some reason, it still gives me the actual value from the SSM store

import os

import boto3
from moto import mock_ssm
import pytest

from handler import lambda_handler

@pytest.fixture
def aws_credentials():
    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
    os.environ["AWS_SECURITY_TOKEN"] = "testing"
    os.environ["AWS_SESSION_TOKEN"] = "testing"

@mock_ssm
def test_ssm():
    ssm = boto3.client('ssm')
    ssm.put_parameter(
        Name="/mypath/password",
        Description="A test parameter",
        Value="this is it!",
        Type="SecureString"
    )
    resp = lambda_handler({"param": "/mypath/password"}, "")
    assert resp["Parameter"]["Value"] == "this is it!"

Am I missing something overhear, what should I do to make it work, or is there an alternative way to mock SSM in python.


Solution

  • When you patched ssm via @mock_ssm, you already have instantiated the variable handler.client in handler.py as a real instance of boto3 client, thus it isn't the patched version.

    Solution 1:

    Initialize the client only during the lambda handler so that the patch is already in effect when it is created.

    handler.py

    import boto3
    
    # client = boto3.client('ssm')  # Remove this instantiation
    
    def lambda_handler(event, context):
        client = boto3.client('ssm')  # Move it here
        return client.get_parameter(Name=event["param"], WithDecryption=True)
    
    if __name__ == '__main__':
        print(lambda_handler({"param": "/mypath/password"}, ""))
    

    Solution 2:

    Only import/execute the handler.py file once the patch is already in effect.

    test_handler.py

    ...
    # from handler import lambda_handler  # Remove this import
    ...
    @mock_ssm
    def test_ssm():
        from handler import lambda_handler  # Move it here
        ...
    ...
    

    Solution 3

    Reload handler.py once the patch is in effect.

    ...
    @mock_ssm
    def test_ssm():
        from importlib import reload
        import sys
        reload(sys.modules['handler'])
        ...
    ...
    

    Output

    $ pytest -q -rP
    ============================================================================================ PASSES ============================================================================================
    1 passed in 1.35s