如何在unittest python中模拟postgres select查询数据

import unittest
from unittest import mock

def get_data_from_postgres():
    # 假设这是一个从PostgreSQL中查询数据的函数
    # 实际情况下，这里会有真正的数据库连接和查询操作
    # 这里我们只返回一个示例的查询结果
    return [{'id': 1, 'name': 'Alice'}, {'id': 2, 'name': 'Bob'}]

def process_data():
    data = get_data_from_postgres()
    # 对查询结果进行处理的函数
    # 这里我们只打印查询结果
    for row in data:
        print(row)

class TestProcessData(unittest.TestCase):
    @mock.patch('__main__.get_data_from_postgres')
    def test_process_data(self, mock_get_data):
        # 模拟查询结果
        mock_get_data.return_value = [{'id': 1, 'name': 'Alice'}, {'id': 2, 'name': 'Bob'}]
        
        # 调用被测试的函数
        process_data()
        
        # 断言查询结果是否被正确处理
        self.assertEqual(mock_get_data.call_count, 1)
        self.assertEqual(mock_get_data.return_value, [{'id': 1, 'name': 'Alice'}, {'id': 2, 'name': 'Bob'}])

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