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.

Answer

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