CODEWITHSUNDEEP
pythonmockingpsycopg2
user4038636
user4038636

Reputation:

Python mock psycopg2 from base class

Hi I am having troubles with mocking cursor object from inherited class. Here is some part of my code:

# classa.py
class ClassA(ClassB):
    def __init__(self, params):
        ClassB.__init__(params)
        # other params

    def get_current_attr(self):
        sql = "SELECT max(attr) FROM {};".format(self.table_name)
        self.rs_cursor.execute(sql)
        current_attr = self.rs_cursor.fetchone()[0]
        return current_attr

# classb.py
import psycopg2

class ClassB:
    def __init__(self, params):
        self.redshift_conn = self.open_redshift_connection()
        self.redshift_conn.set_isolation_level(
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
        self.rs_cursor = self.redshift_conn.cursor()

Now I want to mock psycopg2 cursor in get_current_attr() function from ClassA. I tried to do it like this:

@mock.patch("mypackage.classa.classb.psycopg2.connect")
def test_get_current_attr(self, mock_connect):
    mock_con = mock_connect.return_value
    mock_cur = mock_con.cursor.return_value
    mock_cur.fetch_one.return_value = 1
    result = self.gap_load_instance.get_current_batch_id()
    mock_cur.execute.assert_called_with(
        "SELECT max(attr) FROM tablename"
    )

And I am getting:

raise AssertionError('Expected call: %s\nNot called' % (expected,))

Can anyone help me with this?

Upvotes: 1

Views: 3423

Answers (1)

Psykar
Psykar

Reputation: 11

Mocking the entirety of psycopg2 is easiest. I've added in stubs for your missing functions, but something like this should work:

# classa.py
class ClassA(ClassB):
    def __init__(self, params):
        super().__init__(params)
        # other params

    def get_current_attr(self):
        sql = "SELECT max(attr) FROM {};".format(self.table_name)
        self.rs_cursor.execute(sql)
        current_attr = self.rs_cursor.fetchone()[0]
        return current_attr

# classb.py
import psycopg2

class ClassB:
    def __init__(self, params):
        self.redshift_conn = self.open_redshift_connection()
        self.redshift_conn.set_isolation_level(
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
        self.rs_cursor = self.redshift_conn.cursor()
        self.table_name = "atablename"

    def open_redshift_connection(self):
        return psycopg2.connect()

Then

from unittest import TestCase, mock
from classa import ClassA

class Tester(TestCase):
    @mock.patch("classb.psycopg2")
    def test_get_current_attr(self, mock_psycopg2):
        mock_cur = mock_psycopg2.connect().cursor()
        mock_cur.fetch_one.return_value = 1

        a = ClassA({})
        result = a.get_current_attr()

        mock_cur.execute.assert_called_with(
            "SELECT max(attr) FROM {};".format(a.table_name)
        )

Upvotes: 1

Related Questions