unit test a function with pymysql query.

pymysq_query_unittest.py

from unittest import TestCase, mock
import simple

class TestSimple(TestCase):

    @mock.patch('simple.pymysql', autospec=True)
    def test_get-data(self, mock_pymysql):
        mock_cursor = mock.MagicMock()
        test_data = [{'password': 'secret', 'id': 1}]
        mock_cursor.fetchall.return_value = test_data
        mock_pymysql.connect.return_value.cursor.return_value.__enter__.return_value = mock_cursor

        self.assertEqual(test_data, simple.get_user_data())