How to mock psycopg2 connection methods?

Question

I would like to verify that the commit() and close() methods are being called.

I have the following class:

class ClosingConnection:
    def __init__(self, schema_name: str) -> None:
        super().__init__()
        self.schema_name = schema_name

    def __enter__(self):
        try:
            self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
            )
            return self.conn
        except psycopg2.OperationalError:
            pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.conn:
            if not exc_type:
                self.conn.commit()
            self.conn.close()

@mock.patch("db.connection.psycopg2.connect")
def test_connection_exit(self, conn_mock):
    close_mock = conn_mock.close
    commit_mock = conn_mock.commit

    with ClosingConnection("tenant"):
        assert close_mock.call_count == 1
        assert commit_mock.call_count == 1

Unfortunately the test fails because the function().call_count values are always 0.

Answer 1

ok the main issue was I was trying to assert while still in the with block.

this works:

def test_connection_exit(self, mocker): 
    conn_mock = mocker.patch("db.connection.psycopg2.connect")

    with ClosingConnection("tenant"):
        pass

    assert conn_mock.return_value.close.call_count == 1
    assert conn_mock.return_value.commit.call_count == 1