diff --git a/README.rst b/README.rst index 20991c8..80f82c6 100644 --- a/README.rst +++ b/README.rst @@ -178,8 +178,7 @@ When used as a context manager, time is mocked during the ``with`` block: Class Decorator ^^^^^^^^^^^^^^^ -Only ``unittest.TestCase`` subclasses are supported. -When applied as a class decorator to such classes, time is mocked from the start of ``setUpClass()`` to the end of ``tearDownClass()``: +Decorating a ``unittest.TestCase`` subclass will mock time from the start of ``setUpClass()`` to the end of ``tearDownClass()``: .. code-block:: python @@ -195,6 +194,22 @@ When applied as a class decorator to such classes, time is mocked from the start Note this is different to ``unittest.mock.patch()``\'s behaviour, which is to mock only during the test methods. +Any other class can also be decorated - time will be mocked for each method with a name starting with ``test``: + +.. code-block:: python + + import time + import time_machine + + + @time_machine.travel(0.0) + class TestDeepPast: + def test_in_the_deep_past(self): + assert 0.0 < time.time() < 1.0 + + def not_a_test(self): + assert not 0.0 < time.time() < 1.0 + Timezone mocking ^^^^^^^^^^^^^^^^ diff --git a/src/time_machine/__init__.py b/src/time_machine/__init__.py index ad73eca..166bbc7 100644 --- a/src/time_machine/__init__.py +++ b/src/time_machine/__init__.py @@ -263,7 +263,7 @@ def __exit__( self.stop() @overload - def __call__(self, wrapped: Type[TestCase]) -> Type[TestCase]: # pragma: no cover + def __call__(self, wrapped: type) -> type: # pragma: no cover ... @overload @@ -281,15 +281,11 @@ def __call__( def __call__( self, wrapped: Union[ - Type[TestCase], + type, Callable[..., Coroutine[Any, Any, Any]], Callable[..., Any], ], - ) -> Union[ - Type[TestCase], - Callable[..., Coroutine[Any, Any, Any]], - Callable[..., Any], - ]: + ) -> Union[type, Callable[..., Coroutine[Any, Any, Any]], Callable[..., Any]]: if isinstance(wrapped, type): return self._decorate_class(wrapped) elif inspect.iscoroutinefunction(wrapped): @@ -298,32 +294,44 @@ def __call__( assert callable(wrapped) return self._decorate_callable(wrapped) - def _decorate_class(self, wrapped: Type[TestCase]) -> Type[TestCase]: - if not issubclass(wrapped, TestCase): - raise TypeError("Can only decorate unittest.TestCase subclasses.") + def _decorate_class(self, wrapped: type) -> type: + if issubclass(wrapped, TestCase): + # Modify the setUpClass method + orig_setUpClass = wrapped.setUpClass - # Modify the setUpClass method - orig_setUpClass = wrapped.setUpClass + @functools.wraps(orig_setUpClass) + def setUpClass(cls: Type[TestCase]) -> None: + self.__enter__() + try: + orig_setUpClass() + except Exception: + self.__exit__(*sys.exc_info()) + raise - @functools.wraps(orig_setUpClass) - def setUpClass(cls: Type[TestCase]) -> None: - self.__enter__() - try: - orig_setUpClass() - except Exception: - self.__exit__(*sys.exc_info()) - raise + wrapped.setUpClass = classmethod(setUpClass) # type: ignore[assignment] - wrapped.setUpClass = classmethod(setUpClass) # type: ignore[assignment] + orig_tearDownClass = wrapped.tearDownClass - orig_tearDownClass = wrapped.tearDownClass + @functools.wraps(orig_tearDownClass) + def tearDownClass(cls: Type[TestCase]) -> None: + orig_tearDownClass() + self.__exit__(None, None, None) - @functools.wraps(orig_tearDownClass) - def tearDownClass(cls: Type[TestCase]) -> None: - orig_tearDownClass() - self.__exit__(None, None, None) + wrapped.tearDownClass = classmethod( # type: ignore[assignment] + tearDownClass + ) + else: + # Wrap all methods starting with "test", intended for pytest-like classes + for (attribute, value) in wrapped.__dict__.items(): + if not attribute.startswith("test"): + continue + + if not callable(value) or inspect.isclass(value): + continue + + wrapped_method = self._decorate_callable(value) + setattr(wrapped, attribute, wrapped_method) - wrapped.tearDownClass = classmethod(tearDownClass) # type: ignore[assignment] return wrapped def _decorate_coroutine( diff --git a/tests/test_time_machine.py b/tests/test_time_machine.py index a23a99e..1c7a726 100644 --- a/tests/test_time_machine.py +++ b/tests/test_time_machine.py @@ -487,14 +487,35 @@ async def record_time(): assert recorded_time == EPOCH + 140.0 -def test_class_decorator_fails_non_testcase(): - with pytest.raises(TypeError) as excinfo: +@time_machine.travel(EPOCH) +def test_non_unittest_class_decorator(): + @time_machine.travel(EPOCH + 25.0) + class NonUnitTestClassTests: + def test_class_decorator(self): + assert time.time() == EPOCH + 25.0 - @time_machine.travel(EPOCH) - class Something: - pass + NonUnitTestClassTests().test_class_decorator() + + +@time_machine.travel(EPOCH) +def test_non_unittest_class_decorator_ignores_non_tests(): + @time_machine.travel(EPOCH + 25.0) + class NonUnitTestClassTests: + def not_a_test(self): + assert time.time() == EPOCH + + NonUnitTestClassTests().not_a_test() + + +@time_machine.travel(EPOCH) +def test_non_unittest_class_decorator_ignores_nested_classes(): + @time_machine.travel(EPOCH + 25.0) + class NonUnitTestClassTests: + class test_class: + def test(self): + assert time.time() == EPOCH - assert excinfo.value.args == ("Can only decorate unittest.TestCase subclasses.",) + NonUnitTestClassTests().test_class().test() class MethodDecoratorTests: