From 0b4f34e1279b09c4d4885f14d4516fa9a75fa9da Mon Sep 17 00:00:00 2001 From: "Nathaniel J. Smith" Date: Wed, 4 Mar 2020 17:35:13 -0800 Subject: [PATCH] Fix pendulum.test to properly unwind after an exception Previously, if an exception occurred inside a 'with pendulum.test(...)' block, then the monkeypatch would remain in place instead of being un-done. This could cause confusing results, as one failing test could cause other tests to run with an unexpected mock in place and cause other failures. --- pendulum/helpers.py | 8 ++++---- tests/test_helpers.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) 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