Merge pull request #626 from sdispater/feature/travel

Add time travel for testing

Author: sdispater

pendulum/__init__.py

@@ -29,16 +29,13 @@
 from pendulum.formatting import Formatter
 from pendulum.helpers import format_diff
 from pendulum.helpers import get_locale
-from pendulum.helpers import get_test_now
-from pendulum.helpers import has_test_now
 from pendulum.helpers import locale
 from pendulum.helpers import set_locale
-from pendulum.helpers import set_test_now
-from pendulum.helpers import test
 from pendulum.helpers import week_ends_at
 from pendulum.helpers import week_starts_at
 from pendulum.parser import parse
 from pendulum.period import Period
+from pendulum.testing.traveller import Traveller
 from pendulum.time import Time
 from pendulum.tz import UTC
 from pendulum.tz import local_timezone
@@ -239,8 +236,7 @@ def from_format(
     """
     Creates a DateTime instance from a specific format.
     """
-    parts = _formatter.parse(string, fmt, now(), locale=locale)
-
+    parts = _formatter.parse(string, fmt, now(tz=tz), locale=locale)
     if parts["tz"] is None:
         parts["tz"] = tz
 
@@ -297,6 +293,15 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     return Period(start, end, absolute=absolute)
 
 
+# Testing
+
+_traveller = Traveller(DateTime)
+
+freeze = _traveller.freeze
+travel = _traveller.travel
+travel_to = _traveller.travel_to
+travel_back = _traveller.travel_back
+
 __all__ = [
     "__version__",
     "DAYS_PER_WEEK",
@@ -324,20 +329,17 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     "datetime",
     "duration",
     "format_diff",
+    "freeze",
     "from_format",
     "from_timestamp",
     "get_locale",
-    "get_test_now",
-    "has_test_now",
     "instance",
     "local",
     "locale",
     "naive",
     "now",
     "period",
     "set_locale",
-    "set_test_now",
-    "test",
     "week_ends_at",
     "week_starts_at",
     "parse",
@@ -352,6 +354,9 @@ def period(start: DateTime, end: DateTime, absolute: bool = False) -> Period:
     "timezones",
     "today",
     "tomorrow",
+    "travel",
+    "travel_back",
+    "travel_to",
     "FixedTimezone",
     "Timezone",
     "yesterday",

pendulum/date.py

@@ -24,8 +24,6 @@
 from pendulum.constants import YEARS_PER_DECADE
 from pendulum.exceptions import PendulumException
 from pendulum.helpers import add_duration
-from pendulum.helpers import get_test_now
-from pendulum.helpers import has_test_now
 from pendulum.mixins.default import FormattableMixin
 from pendulum.period import Period
 
@@ -733,9 +731,6 @@ def average(self, dt: date | None = None) -> Date:
 
     @classmethod
     def today(cls) -> Date:
-        if has_test_now():
-            return cast(pendulum.DateTime, get_test_now()).date()
-
         dt = date.today()
 
         return cls(dt.year, dt.month, dt.day)

pendulum/datetime.py

@@ -32,8 +32,6 @@
 from pendulum.date import Date
 from pendulum.exceptions import PendulumException
 from pendulum.helpers import add_duration
-from pendulum.helpers import get_test_now
-from pendulum.helpers import has_test_now
 from pendulum.period import Period
 from pendulum.time import Time
 from pendulum.tz import UTC
@@ -135,15 +133,6 @@ def now(
         """
         Get a DateTime instance for the current date and time.
         """
-        if has_test_now():
-            test_instance: DateTime = cast(DateTime, get_test_now())
-            _tz = pendulum._safe_timezone(tz)
-
-            if tz is not None and _tz != test_instance.timezone:
-                test_instance = test_instance.in_tz(_tz)
-
-            return test_instance
-
         if tz is None or tz == "local":
             dt = datetime.datetime.now(local_timezone())
         elif tz is UTC or tz == "UTC":

pendulum/helpers.py

@@ -3,13 +3,11 @@
 import os
 import struct
 
-from contextlib import contextmanager
 from datetime import date
 from datetime import datetime
 from datetime import timedelta
 from math import copysign
 from typing import TYPE_CHECKING
-from typing import Iterator
 from typing import TypeVar
 from typing import overload
 
@@ -178,27 +176,6 @@ def _sign(x: float) -> int:
 # Global helpers
 
 
-@contextmanager
-def test(mock: pendulum.DateTime) -> Iterator[None]:
-    set_test_now(mock)
-    try:
-        yield
-    finally:
-        set_test_now()
-
-
-def set_test_now(test_now: pendulum.DateTime | None = None) -> None:
-    pendulum._TEST_NOW = test_now
-
-
-def get_test_now() -> pendulum.DateTime | None:
-    return pendulum._TEST_NOW
-
-
-def has_test_now() -> bool:
-    return pendulum._TEST_NOW is not None
-
-
 def locale(name: str) -> Locale:
     return Locale.load(name)
 
@@ -238,10 +215,6 @@ def week_ends_at(wday: int) -> None:
     "week_day",
     "add_duration",
     "format_diff",
-    "test",
-    "set_test_now",
-    "get_test_now",
-    "has_test_now",
     "locale",
     "set_locale",
     "get_locale",

pendulum/parser.py

@@ -24,7 +24,7 @@
 
 def parse(text: str, **options: t.Any) -> Date | Time | DateTime | Duration:
     # Use the mock now value if it exists
-    options["now"] = options.get("now", pendulum.get_test_now())
+    options["now"] = options.get("now")
 
     return _parse(text, **options)
 

pendulum/testing/__init__.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import cast
+
+from pendulum.datetime import DateTime
+from pendulum.utils._compat import PYPY
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+
+class BaseTraveller:
+    def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
+        self._datetime_class: type[DateTime] = datetime_class
+
+    def freeze(self: BaseTraveller) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel_back(self: BaseTraveller) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel(
+        self,
+        years: int = 0,
+        months: int = 0,
+        weeks: int = 0,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: int = 0,
+        microseconds: int = 0,
+    ) -> BaseTraveller:
+        raise NotImplementedError()
+
+    def travel_to(self, dt: DateTime) -> BaseTraveller:
+        raise NotImplementedError()
+
+
+if not PYPY:
+    import time_machine
+
+    class Traveller(BaseTraveller):
+        def __init__(self, datetime_class: type[DateTime] = DateTime) -> None:
+            super().__init__(datetime_class)
+
+            self._started: bool = False
+            self._traveller: time_machine.travel | None = None
+            self._coordinates: time_machine.Coordinates | None = None
+
+        def freeze(self) -> Traveller:
+            if self._started:
+                cast(time_machine.Coordinates, self._coordinates).move_to(
+                    self._datetime_class.now(), tick=False
+                )
+            else:
+                self._start(freeze=True)
+
+            return self
+
+        def travel_back(self) -> Traveller:
+            if not self._started:
+                return self
+
+            cast(time_machine.travel, self._traveller).stop()
+            self._coordinates = None
+            self._traveller = None
+            self._started = False
+
+            return self
+
+        def travel(
+            self,
+            years: int = 0,
+            months: int = 0,
+            weeks: int = 0,
+            days: int = 0,
+            hours: int = 0,
+            minutes: int = 0,
+            seconds: int = 0,
+            microseconds: int = 0,
+            *,
+            freeze: bool = False,
+        ) -> Traveller:
+            self._start(freeze=freeze)
+
+            cast(time_machine.Coordinates, self._coordinates).move_to(
+                self._datetime_class.now().add(
+                    years=years,
+                    months=months,
+                    weeks=weeks,
+                    days=days,
+                    hours=hours,
+                    minutes=minutes,
+                    seconds=seconds,
+                    microseconds=microseconds,
+                )
+            )
+
+            return self
+
+        def travel_to(self, dt: DateTime, *, freeze: bool = False) -> Traveller:
+            self._start(freeze=freeze)
+
+            cast(time_machine.Coordinates, self._coordinates).move_to(dt)
+
+            return self
+
+        def _start(self, freeze: bool = False) -> None:
+            if self._started:
+                return
+
+            if not self._traveller:
+                self._traveller = time_machine.travel(
+                    self._datetime_class.now(), tick=not freeze
+                )
+
+            self._coordinates = self._traveller.start()
+
+            self._started = True
+
+        def __enter__(self) -> Traveller:
+            self._start()
+
+            return self
+
+        def __exit__(
+            self,
+            exc_type: type[BaseException] | None,
+            exc_val: BaseException | None,
+            exc_tb: TracebackType,
+        ) -> None:
+            self.travel_back()
+
+else:
+
+    class Traveller(BaseTraveller):  # type: ignore[no-redef]
+
+        ...