python - 如何模拟 psycopg2 的紧密连接？

Question

我想使用 psycopg2 测试我的数据库类，但我遇到了问题 - 我不知道如何测试 close() 方法。我尝试了一些想法，这是下面的最后一个想法，但它不起作用。所以有人知道如何嘲笑吗？

数据库.py

import psycopg2


class DB:
    def __init__(self, name: str) -> None:
        self.name = name
        self.conn = psycopg2.connect(dbname=self.name,
                                     host="localhost",
                                     port="5432",
                                     user="postgres",
                                     password="postgres")
        
    def disconnect(self):
        self.conn.close()

test_db.py

import unittest
from unittest import mock

import psycopg2

from db import DB


class TestDB(unittest.TestCase):

    def setUp(self) -> None:
        self.db = DB("testdb")

    def tearDown(self) -> None:
        self.db.disconnect()

    @mock.patch("db.psycopg2.connect")
    def test_DB_should_connect_to_db(self, mock_connect):
        DB("testdb")
        mock_connect.assert_called_with(dbname="testdb",
                                        host="localhost",
                                        port="5432",
                                        user="postgres",
                                        password="postgres")

    @mock.patch("db.psycopg2.connect")
    def test_db_should_disconnect_from_db(self, mock_connect):
        mock_conn = mock_connect.return_value
        mock_close = mock_conn.close
        self.db.disconnect()
        mock_close.assert_called_once()





if __name__ == '__main__':
    unittest.main()

标签： pythonpostgresqlpsycopg2