timeboat.py

import datetime
import pytest
from freezegun import freeze_time

@pytest.fixture()
def timeboat():
    """Grab the helm of the timeboat, we're going on an adventure!
    Allows the test writer to freeze time so that time.time() returns a controlled value
    https://stackoverflow.com/a/28073449/4717806"""

    class Timeboat():
        """Timeboat Class"""
        @classmethod
        def freeze_now(cls, tz=None):
            """Freezes time to whatever now is"""
            if tz:
                cls.freezegun_ctx = freeze_time(tz.localize(datetime.datetime.now()))
            else:
                cls.freezegun_ctx = freeze_time(datetime.datetime.now())
            cls.freezegun_ctx.start()

        @classmethod
        def freeze(cls, *args, tz=None, **kwargs):
            """Freezes time at the specified value"""
            if tz:
                cls.freezegun_ctx = freeze_time(tz.localize(datetime.datetime(*args, **kwargs)))
            else:
                cls.freezegun_ctx = freeze_time(datetime.datetime(*args, **kwargs))
            cls.freezegun_ctx.start()

        @classmethod
        def unfreeze(cls):
            """Unfreezes time"""
            if hasattr(cls, 'freezegun_ctx'):
                cls.freezegun_ctx.stop()
                del cls.freezegun_ctx

        @classmethod
        def shift(cls, timedelta=None, **kwargs):
            """ Moves time fwd/bwd by the delta"""
            cls.freezegun_ctx.stop()
            if not timedelta:
                timedelta = datetime.timedelta(**kwargs)
            cls.freezegun_ctx.time_to_freeze += timedelta
            cls.freezegun_ctx.start()

    yield Timeboat
    Timeboat.unfreeze()