Allow set_testing to be used as a context manager

CHANGELOG.md

@@ -22,6 +22,9 @@ You can find our backwards-compatibility policy [here](https://github.com/hynek/
   When `cap=True`, it is used as an upper cap; that means that if the original attempts number is lower, it's not changed.
   [#80](https://github.com/hynek/stamina/pull/80)
 
+- `stamina.set_testing()` can now be used as a context manager.
+  [#94](https://github.com/hynek/stamina/pull/94)
+
 
 ## [24.3.0](https://github.com/hynek/stamina/compare/24.2.0...24.3.0) - 2024-08-27
 

src/stamina/_config.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from threading import Lock
+from types import TracebackType
 from typing import Callable
 
 from .instrumentation import RetryHookFactory
@@ -154,9 +155,25 @@ def is_testing() -> bool:
     return CONFIG.testing is not None
 
 
+class _RestoreTestingCM:
+    def __init__(self, old: _Testing | None) -> None:
+        self.old = old
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        CONFIG.testing = self.old
+
+
 def set_testing(
     testing: bool, *, attempts: int = 1, cap: bool = False
-) -> None:
+) -> _RestoreTestingCM:
     """
     Activate or deactivate test mode.
 
@@ -170,5 +187,9 @@ def set_testing(
 
     .. versionadded:: 24.3.0
     .. versionadded:: 25.1.0 *cap*
+    .. versionadded:: 25.1.0 Can be used as a context manager.
     """
+    old = CONFIG.testing
     CONFIG.testing = _Testing(attempts, cap) if testing else None
+
+    return _RestoreTestingCM(old)

tests/test_async.py

@@ -306,3 +306,24 @@ def test_repr(self):
         assert f"<BoundAsyncRetryingCaller(ValueError, {r})>" == repr(
             arc.on(ValueError)
         )
+
+
+async def test_testing_mode_context():
+    """
+    Testing mode context manager works with async code.
+    """
+    assert not stamina.is_testing()
+
+    with stamina.set_testing(True, attempts=3):
+        assert stamina.is_testing()
+
+        with pytest.raises(ValueError):  # noqa: PT012
+            async for attempt in stamina.retry_context(on=ValueError):
+                assert 0.0 == attempt.next_wait
+
+                with attempt:
+                    raise ValueError
+
+        assert 3 == attempt.num
+
+    assert not stamina.is_testing()

tests/test_config.py

@@ -2,10 +2,11 @@
 #
 # SPDX-License-Identifier: MIT
 
+from contextlib import suppress
 from threading import Lock
 
-from stamina import is_active, set_active
-from stamina._config import _Config, _Testing
+from stamina import is_active, is_testing, set_active, set_testing
+from stamina._config import CONFIG, _Config, _Testing
 
 
 def test_activate_deactivate():
@@ -67,3 +68,49 @@ def test_cap_true_with_none(self):
         t = _Testing(100, True)
 
         assert 100 == t.get_attempts(None)
+
+    def test_context_manager(self):
+        """
+        set_testing works as a context manager.
+        """
+        assert not is_testing()
+
+        with set_testing(True, attempts=3):
+            assert is_testing()
+            assert 3 == CONFIG.testing.get_attempts(None)
+            assert not CONFIG.testing.cap
+
+        assert not is_testing()
+
+    def test_context_manager_nested(self):
+        """
+        set_testing context managers can be nested.
+        """
+        assert not is_testing()
+
+        with set_testing(True, attempts=3):
+            assert is_testing()
+            assert CONFIG.testing.attempts == 3
+
+            with set_testing(True, attempts=5, cap=True):
+                assert is_testing()
+                assert CONFIG.testing.attempts == 5
+                assert CONFIG.testing.cap
+
+            assert is_testing()
+            assert CONFIG.testing.attempts == 3
+            assert not CONFIG.testing.cap
+
+        assert not is_testing()
+
+    def test_context_manager_exception(self):
+        """
+        set_testing context manager restores state even if an exception occurs.
+        """
+        assert not is_testing()
+
+        with suppress(ValueError), set_testing(True, attempts=3):
+            assert is_testing()
+            raise ValueError("test")
+
+        assert not is_testing()

tests/test_sync.py

@@ -290,3 +290,24 @@ def test_repr(self):
         assert f"<BoundRetryingCaller(ValueError, {r})>" == repr(
             rc.on(ValueError)
         )
+
+
+def test_testing_mode_context():
+    """
+    Testing mode context manager works with sync code.
+    """
+    assert not stamina.is_testing()
+
+    with stamina.set_testing(True, attempts=3):
+        assert stamina.is_testing()
+
+        with pytest.raises(ValueError):  # noqa: PT012
+            for attempt in stamina.retry_context(on=ValueError):
+                assert 0.0 == attempt.next_wait
+
+                with attempt:
+                    raise ValueError
+
+        assert 3 == attempt.num
+
+    assert not stamina.is_testing()