使用Pytest测试psycopg的execute_query方法断言失败求助

我尝试测试PostgresqlConnection类的execute_query方法，但测试无法正常运行，出现如下断言错误：

AssertionError: assert [] == [{'id': 1, 'n...lary': 70000}]
E Right contains 3 more items, first extra item: {'id': 1,
'name': 'Alice', 'salary': 50000}
E Use -v to get more diff

相关实现代码 #

class PostgresqlConnection():
    dbname: str = PG_ETL_DATABASE
    user: str = PG_ETL_USER
    password: str = PG_ETL_PASSWORD
    host: str = PG_ETL_HOST
    port: str = PG_ETL_PORT

    @property
    def connection(self):
        return psycopg.connect(**self.dict())

    def execute_query(self, query, is_fetch_result=False, is_commit=False):
        """Method execute a query on the database"""
        with self.connection as conn:
            with conn.cursor() as cursor:
                cursor.execute(query)
                results = []
                if is_fetch_result:
                    print(f"cursor.description1: {cursor.description.return_value}")
                    column_names = [desc[0] for desc in cursor.description]
                    print(f"column_names: {column_names}")
                    rows = cursor.fetchall()
                    print(f"rows: {rows}")
                    for row in rows:
                        results.append(dict(zip(column_names, row)))
            if is_commit:
                conn.commit()
            return results

测试代码 #

@mock.patch("psycopg.connect")
def test_execute_query(self, mock_connect):
    """test execute query"""
    expected = [
        {"id": 1, "name": "Alice", "salary": 50000},
        {"id": 2, "name": "Bob", "salary": 60000},
        {"id": 3, "name": "Charlie", "salary": 70000},
    ]

    mock_con = mock_connect.return_value
    mock_cur = mock_con.cursor.return_value
    mock_cur.description.return_value = [
        ("id", None, None, None, None, None, None),
        ("name", None, None, None, None, None, None),
        ("salary", None, None, None, None, None, None),
    ]
    mock_cur.fetchall.return_value = [
        (1, "Alice", 50000),
        (2, "Bob", 60000),
        (3, "Charlie", 70000),
    ]
    result = self.resource.execute_query(
        query="q", is_fetch_result=True, is_commit=False
    )
    assert result == expected

问题原因及修复方案 #

1. 实现代码错误 #

cursor.description是属性而非方法，不需要调用return_value。原代码中print(f"cursor.description1: {cursor.description.return_value}")会导致无法正确获取列名，进而无法生成结果字典。

修正该行：

print(f"cursor.description1: {cursor.description}")

2. 测试代码mock错误 #

同样，cursor.description是属性，测试中不应设置return_value，直接给mock_cur.description赋值即可。

修正测试代码中的mock设置：

mock_cur.description = [
    ("id", None, None, None, None, None, None),
    ("name", None, None, None, None, None, None),
    ("salary", None, None, None, None, None, None),
]

完成以上两处修正后，测试即可正常运行，断言会通过。

内容的提问来源于stack exchange，提问作者idan ovadia