diff --git a/pendulum/helpers.py b/pendulum/helpers.py index 5085a38b..df9aa8c9 100644 --- a/pendulum/helpers.py +++ b/pendulum/helpers.py @@ -139,10 +139,10 @@ def _sign(x): @contextmanager def test(mock): set_test_now(mock) - - yield - - set_test_now() + try: + yield + finally: + set_test_now() def set_test_now(test_now=None): diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 95f90beb..37dc4185 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -190,3 +190,20 @@ def test_week_ends_at_invalid_value(): with pytest.raises(ValueError): pendulum.week_ends_at(11) + + +def test_with_test(): + t = pendulum.datetime(2000, 1, 1) + + with pendulum.test(t): + assert pendulum.now() == t + + assert pendulum.now() != t + + # Also make sure that it restores things after an exception + with pytest.raises(RuntimeError): + with pendulum.test(t): + assert pendulum.now() == t + raise RuntimeError + + assert pendulum.now() != t