Python Mocking AWS Wrangler / DB Connection

49 Views Asked by At

I'm new to Python and testing and trying to write some tests that interact with AWS wrangler in my Python 3 lambda function. Here is the documentation.

The function I want to test looks like

def finalize_request(request_id: str):
    logging.info(f"Marking request: {request_id} as complete")

    try:
        with wr.postgresql.connect(DB_CONNECTION_NAME) as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    f"UPDATE log.account_status_chart_request SET processed_date = NOW() WHERE request_id = '{request_id}'"
                )
                conn.commit()
    except Exception as e:
        logging.error(f"An exception occurred while finalizing the request: {e}")
        raise e

I've sucessfully intercepted the wr.postgresql.connect portion of the code using

def test_finalize_request(mocker):
    mock_connection = mocker.MagicMock()

    @contextmanager
    def mock_connect(*args, **kwargs):
        yield mock_connection

    finalize_request(request_id=1)

What I'm trying to do is within my mock_connection specifically grap the cursor.execute and override it. I'm pushing in some dummy data to a sqlite database and want to do my update in that database.

Thoughts?

The expectation is that I have another function called mock_update or something of the likes that simply runs something like

connection.execute(
        <sql parameter extracted from the original cursor.execute>
    )
connection.commit()

I can then read the value to make sure the processed date is set.

0

There are 0 best solutions below