Python 如何模拟单个单元测试的本地时区?

Python 如何模拟单个单元测试的本地时区?,python,timezone,python-unittest,python-mock,Python,Timezone,Python Unittest,Python Mock,我正在为python数据库集成编写单元测试。当本地时区与数据库时区不同时,我希望确保正确处理时区。请参见下面的示例 我想在一次测试中设置本地时区。我有一个可行的解决方案,但怀疑有更好的方法来达到同样的效果 这是我的密码: import datetime import os import time import unittest from functools import wraps import pytz class MyFakeDatabase: ""&quo

我正在为python数据库集成编写单元测试。当本地时区与数据库时区不同时,我希望确保正确处理时区。请参见下面的示例

我想在一次测试中设置本地时区。我有一个可行的解决方案,但怀疑有更好的方法来达到同样的效果

这是我的密码:

import datetime
import os
import time
import unittest
from functools import wraps

import pytz


class MyFakeDatabase:
    """
    This imitates a database that stores timestamps in UTC
    """

    def __init__(self):
        self.records = []

    @staticmethod
    def _get_timestamp_in_utc(timestamp):
        return timestamp.astimezone(pytz.utc).replace(tzinfo=None)

    def add_record(self, title, expiry_date):
        new_record = {
            "id": len(self.records),
            "title": title,
            "expiry_date": self._get_timestamp_in_utc(expiry_date),
        }
        self.records.append(new_record)
        return new_record["id"]

    def get_active_records(self):
        # now = datetime.datetime.utcnow() # will pass
        now = (
            datetime.datetime.now()
        )  # will fail as this gets a naive datetime based on the local time, but the database stores them as UTC - this is the case I am writing tests for
        return list(filter(lambda record: record["expiry_date"] > now, self.records))


def set_local_timezone(timezone):
    def set_timezone(new_timezone):
        os.environ["TZ"] = new_timezone
        time.tzset()

    def decorator(func):
        @wraps(func)
        def timezone_adjusted_func(*args, **kwargs):
            current_local_timezone = str(datetime.datetime.now().astimezone().tzinfo)
            result = None
            try:
                set_timezone(timezone)
                result = func(*args, **kwargs)
            finally:
                set_timezone(current_local_timezone)
            return result

        return timezone_adjusted_func

    return decorator


class TestCase(unittest.TestCase):
    @set_local_timezone("America/Chicago")
    def test_get_non_expired_records_respects_timezone(
        self,
    ):
        database = MyFakeDatabase()
        timezone = pytz.timezone("America/Chicago")
        database.add_record(
            "Expired record",
            datetime.datetime.now(timezone) - datetime.timedelta(minutes=5),
        )
        active_id = database.add_record(
            "Active record",
            datetime.datetime.now(timezone) + datetime.timedelta(minutes=5),
        )

        records = database.get_active_records()

        self.assertEqual([active_id], [record["id"] for record in records])

我曾考虑尝试
freezegun
freeze\u time
,但它为
datetime.now()
datetime.utcnow()
返回相同的值,这与我的实现冲突(请参阅)

有更好的方法吗