如何模拟psycopg2连接方法？

Question

我想验证是否正在调用commit（）和close（）方法

我有以下课程：

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

不幸的是，测试失败，因为函数（）的调用计数值始终为0

Answer 1

好的，主要的问题是我仍然在with块中试图断言

这项工作：

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