Pytest测试：如何Mock含日期过滤的函数内数据表的datetime列？

我有一个用于创建Pandas DataFrame的函数，其中包含一个表示实例entry_date的datetime列。该函数在构建DataFrame时，会通过过滤条件仅保留过去两周的entry_date数据。我已准备好测试用的Mock文件，但当测试数据的entry_date超出两周范围时，调用函数后得到的DataFrame为空。我尝试用freeze_time冻结时间，但无法解决问题。相关代码如下：

def create_dataframe(Inherit_from_parent):
    tb1 = get_table("People")
    tb2 = get_table("Record")
    query = sa.select(tb1.c.id, tb1.c.name, tb1.c.entry_date, ... 
            from .... where tb1.entry_date >= date_today() - text('interval "2 weeks"'))


@pytest.fixture
def engine_factory(postgresql):
    def _get_engine(prefix):
        engine = get_database_engine(
            postgresql.info.dbname,
            username=postgresql.info.user,
            password=postgresql.info.password,
            host=postgresql.info.host,
            port=postgresql.info.port,
    )

    tables = {prefix + "_records": "record", prefix + "_people": "people"}

    for filename, tablename in tables.items():
        df = pd.read_parquet(f".{PATH}{filename}.parquet")
        df.to_sql(tablename, engine, index=False)

    return engine

return _get_engine

@freeze_time("2024-04-23")
def test_create_dataframe():
    people_ids = [1,2]

    df = create_dataframe(Parent_class_info(people_ids,**kwargs))
    df_expected = pd.read_parquet(f'{PATH}/df_expected.parquet')

    pd.testing.assert_frame_equal(df, df_expected)

目前仅当测试数据的entry_date为当前日期且配合freeze_time时测试通过，一旦日期过早测试就失败。请问如何在该场景下Mock函数内创建的数据表的datetime列？

方法1：动态调整测试数据的entry_date #

在加载测试Parquet文件后，根据冻结的时间动态修改entry_date，确保数据落在过滤范围内：

import datetime
import pandas as pd

@pytest.fixture
def engine_factory(postgresql):
    def _get_engine(prefix):
        engine = get_database_engine(
            postgresql.info.dbname,
            username=postgresql.info.user,
            password=postgresql.info.password,
            host=postgresql.info.host,
            port=postgresql.info.port,
        )

        tables = {prefix + "_records": "record", prefix + "_people": "people"}
        # 获取冻结的当前日期
        frozen_today = datetime.date.today()

        for filename, tablename in tables.items():
            df = pd.read_parquet(f".{PATH}{filename}.parquet")
            # 调整entry_date到冻结日期的两周范围内
            if 'entry_date' in df.columns:
                # 生成从两周前到今天的日期序列，匹配数据行数
                df['entry_date'] = pd.date_range(
                    start=frozen_today - datetime.timedelta(weeks=2),
                    end=frozen_today,
                    periods=len(df)
                )
            df.to_sql(tablename, engine, index=False)

        return engine
    return _get_engine

不管原Parquet文件里的entry_date是什么，加载后都会被调整到冻结时间的有效范围内，保证过滤条件能命中数据。

方法2：Mock查询中的日期计算逻辑 #

直接替换date_today()函数，让它返回你指定的日期，完全控制过滤条件的时间范围：

from unittest.mock import patch
import datetime

@freeze_time("2024-04-23")
def test_create_dataframe():
    people_ids = [1,2]
    # 自定义目标日期，让两周前刚好覆盖测试数据的最早entry_date
    target_today = datetime.date(2024, 4, 23)
    
    with patch('your_module.date_today') as mock_date_today:
        mock_date_today.return_value = target_today
        df = create_dataframe(Parent_class_info(people_ids,**kwargs))
        df_expected = pd.read_parquet(f'{PATH}/df_expected.parquet')
        pd.testing.assert_frame_equal(df, df_expected)

注意替换your_module为date_today()所在的实际模块名。只要调整target_today，就能让测试数据的entry_date落在过滤范围内。

方法3：直接Mock查询逻辑（单元测试场景） #

如果不需要实际数据库操作，直接Mock查询相关的函数，返回预期的DataFrame：

from unittest.mock import patch, MagicMock
import pandas as pd

def test_create_dataframe():
    people_ids = [1,2]
    # 加载预期结果
    mock_df = pd.read_parquet(f'{PATH}/df_expected.parquet')
    
    with patch('your_module.get_table') as mock_get_table, \
         patch('your_module.sa.select') as mock_select:
        # 模拟表对象
        mock_tb1 = MagicMock()
        mock_tb1.c.id = MagicMock()
        mock_tb1.c.name = MagicMock()
        mock_tb1.c.entry_date = MagicMock()
        mock_get_table.side_effect = [mock_tb1, MagicMock()]
        
        # 模拟查询执行结果
        mock_query = MagicMock()
        mock_query.execute.return_value.fetch_dataframe.return_value = mock_df
        mock_select.return_value = mock_query
        
        df = create_dataframe(Parent_class_info(people_ids,**kwargs))
        pd.testing.assert_frame_equal(df, mock_df)

这种方式完全绕过数据库，直接验证函数逻辑，适合快速单元测试。

内容的提问来源于stack exchange，提问作者Craciun Bogdan