From 8d8216977b1f1509b7cb410756a2dfd25f5a1520 Mon Sep 17 00:00:00 2001 From: Hynek Schlawack Date: Sat, 22 Feb 2025 12:50:06 +0100 Subject: [PATCH 1/2] Allow set_testing to be used as a context manager implements https://github.com/hynek/stamina/discussions/85 --- CHANGELOG.md | 2 ++ src/stamina/_config.py | 23 ++++++++++++++++++- tests/test_async.py | 21 +++++++++++++++++ tests/test_config.py | 51 ++++++++++++++++++++++++++++++++++++++++-- tests/test_sync.py | 21 +++++++++++++++++ 5 files changed, 115 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f08203..d7c96cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,8 @@ 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. + ## [24.3.0](https://github.com/hynek/stamina/compare/24.2.0...24.3.0) - 2024-08-27 diff --git a/src/stamina/_config.py b/src/stamina/_config.py index 2c8e550..cdbdc18 100644 --- a/src/stamina/_config.py +++ b/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) diff --git a/tests/test_async.py b/tests/test_async.py index 1acf413..d0dfd6b 100644 --- a/tests/test_async.py +++ b/tests/test_async.py @@ -306,3 +306,24 @@ def test_repr(self): assert f"" == 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() diff --git a/tests/test_config.py b/tests/test_config.py index 528421f..d9d3fec 100644 --- a/tests/test_config.py +++ b/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() diff --git a/tests/test_sync.py b/tests/test_sync.py index fc996fa..f4bbfc3 100644 --- a/tests/test_sync.py +++ b/tests/test_sync.py @@ -290,3 +290,24 @@ def test_repr(self): assert f"" == 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() From a648901267efc6897dfb8fa62f9fd387b0890adb Mon Sep 17 00:00:00 2001 From: Hynek Schlawack Date: Sat, 22 Feb 2025 12:56:24 +0100 Subject: [PATCH 2/2] Add PR link --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d7c96cb..5fa02c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ You can find our backwards-compatibility policy [here](https://github.com/hynek/ [#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