freeze.py

from google.appengine.api import datastore_types
from mock import patch
from freezegun import freeze_time as _freeze_time
from freezegun.api import FakeDatetime


class freeze_time(object):
    def __init__(self, *args, **kwargs):
        self._gae_patch = patch(
            'google.appengine.ext.db.DateTimeProperty.data_type',
            new=FakeDatetime)
        self._freeze = _freeze_time(*args, **kwargs)
        datastore_types._VALIDATE_PROPERTY_VALUES[FakeDatetime] = \
            datastore_types.ValidatePropertyNothing
        datastore_types._PACK_PROPERTY_VALUES[FakeDatetime] = \
            datastore_types.PackDatetime
        datastore_types._PROPERTY_MEANINGS[FakeDatetime] = \
            datastore_types.entity_pb.Property.GD_WHEN

    def start(self):
        self._freeze.start()
        self._gae_patch.start()

    def stop(self):
        self._freeze.stop()
        self._gae_patch.stop()

    def __enter__(self):
        return self.start()

    def __exit__(self, *args):
        self.stop()