Блог

Я хотел бы убедиться, что вызываются методы 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
 

К сожалению, тест завершается неудачей, потому что значения function().call_count всегда равны 0.

хорошо, главная проблема заключалась в том, что я пытался утверждать, все еще находясь в 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