Mocking pymysql.connect calls in Python unit test

  database, mocking, pymysql, python, unit-testing

I have a function that looks like this that I’m trying to mock.

code.app.py

import pymysql

# conn returns a pymysql.connect connection
conn = get_conn(db_port, db_username, db_name)


def does_table_exist(conn, table_name):
    
    sql = f"SELECT * FROM information_schema.tables WHERE table_name='{table_name}';"

    try:
        table_exists = None
        with conn.cursor() as cur:
            sql_result = cur.execute(sql)
            
            if sql_result == 1:
                table_exists = True
            if sql_result == 0:
                table_exists = False

            return table_exists
    except Exception as error:
        logger.error(error)
        raise
    finally:
        cur.close()

My test file looks like as follows.

test_code.py

def test_check_table_exists(init):
    with patch('code.app.pymysql.connect') as mock_connect:
        cursor = MagicMock()
        execute = Mock(return_value=1)
        cursor.return_value = execute
        mock_connect.return_value.__enter__.return_value = cursor
        
        # Here init is a pytest fixture that imports the module
        response = init.does_table_exist(mock_connect, 'test')
        assert response == True

When I run the above, I get the following

FAILED tests/unit/test_code.py::test_does_table_exist – assert None == True

Mocks totally baffle me at times. I need to ensure that cur.execute is mocked to return 1 back. However None being returned by the function.
Can someone point me to how to properly mock this.

Source: Python Questions

LEAVE A COMMENT