Mocking a module selectively in Python 3.7

Question

I am trying to write unit test for one of my module using pymysql as follows:

def execute(self, query):

    try:
        conn = pymysql.connect(host='localhost', port=3306, user='root', passwd='', 
                                db='mysql')

        cur = conn.cursor()

        cur.execute(query)

    except pymysql.Error as e:
        raise
    except Exception as e:
        raise

While writing unit test for the above function, I am mocking pysql as follows @patch(my_package.my_module.pymysql). As a result pymysql.erroris also becoming mocked. So while testing the failure scenario as follows:

@patch('my_package.my_module.pymysql')
def test_execute_with_failure(self, mock_pymysql):
    ...
    self.my_obj.cur.execute.side_effect = pymysql.Error
    with self.assertRaises(pymysql.Error) as context:
        _ = self.my_obj.execute(query="SELECT FOO FROM BAR")

I am getting the below error:

TypeError: catching classes that do not inherit from BaseException is not allowed

In this regard I have gone through this: Can't catch mocked exception because it doesn't inherit BaseException. But I am not getting how to make this work for pymysql.

EDIT 1: As suggested, I have tried to mock only pymysql.connect as follows:

@patch('my_package.my_module.pymysql.connect')
def test_execute_with_failure(self, mock_pymysql_conn):
     cursor = MagicMock()

     mock_pymysql_conn.return_value.cursor = cursor
     self.my_obj.cur = cursor

     self.my_obj.cur.execute.side_effect = pymysql.Error

     with self.assertRaises(pymysql.Error) as context:
        _ = self.my_obj.execute(query="SELECT FOO FROM BAR")

But, I am getting here, as follows:

E           AttributeError: <module 'pymysql' from '/my_project/venv/lib/python3.7/site-packages/pymysql/__init__.py'> does not have the attribute ''

EDIT 2: I have tried with below approach as well:

@patch('my_package.my_module.pymysql')
def test_execute_with_failure(self, mock_pymysql):
    conn = Mock()
    mock_pymysql.connect.return_value = conn
    cursor = MagicMock()

    conn.return_value.cursor = cursor
    self.my_obj.cur = cursor

    mock_pymysql.Error = type('PymysqlError', (Exception,), {})
    self.my_obj.cur.execute.side_effect = mock_pymysql.Error
    with self.assertRaises(mock_pymysql.Error) as context:
        _ = self.my_obj.execute(query="SELECT * FROM TEST")

Still the same error as Edit 1.

Mock `pymysql.connect` only. I recommend importing the `pymysql` module in the test case as well, then using `@patch.object(pymysql, 'connect')` to mock the `connect` function only from that module. — filbranden, Mar 14 '20 at 06:53
@filbranden thanks a ton for your response. Yeah I have tried the same, but still getting error per edit in OP. — Joy, Mar 14 '20 at 07:04
`mock_pymysql_conn.cursor.return_value` => `mock_pymysql_conn.return_value.cursor`. The `cursor` is an attribute of the conn, which is the return value of the `connection{}` call. — filbranden, Mar 14 '20 at 07:08
@filbranden I am sorry. Still getting the same error. I have edited the OP. May be, I am missing something. — Joy, Mar 14 '20 at 07:16
Hmmm it's possible that this is a module implemented in C, then you can't mock individual functions... Then try again mocking the whole module, but then set `mock_pymysql.Error = type('PymysqlError', (Exception,), {})` to make sure that that name is still an exception in the mock... — filbranden, Mar 14 '20 at 07:24
@filbranden Still getting the same error. Edited the OP accordingly. — Joy, Mar 14 '20 at 07:37
The last one should have `conn.cursor.return_value = cursor`. In fact, that was probably the issue with the second one too... — filbranden, Mar 14 '20 at 07:57
Yeah second one needed `mock_pymysql_conn.return_value.cursor.return_value`. since you call.`connect()` and then call `cursor()` on the thing that returned. — filbranden, Mar 14 '20 at 07:59

Mocking a module selectively in Python 3.7

0 Answers0