Mocking pymysql.connect calls in Python unit test

Question

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.

Answer 1

I was able to get this to work in the following way.

@pytest.mark.parametrize('data, output', [
    (1, True), (0, False)
])
def test_does_table_exist(init, data, output):
    with patch('code.app.pymysql.connect') as mock_connect:
        cursor = MagicMock()
        cursor.execute.return_value = data
        mock_connect.return_value.cursor.return_value.__enter__.return_value = cursor
        t = pymysql.connect()
        response = init.does_table_exist(t, 'test')
        assert response == output