From 3fbd386e6c576ef8edd48b480e3c84a6ef6feb55 Mon Sep 17 00:00:00 2001
From: Pavel Bitiukov <bp72@users.noreply.github.com>
Date: Fri, 1 Mar 2024 23:18:54 +0000
Subject: [PATCH] Fix class decorator for classmethod overrides (#404)

Co-authored-by: Adam Johnson <me@adamj.eu>
---
 CHANGELOG.rst                |  4 ++++
 src/time_machine/__init__.py | 10 ++++++----
 tests/test_time_machine.py   | 24 ++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index c6afce8..c01872b 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,6 +6,10 @@ Changelog
 
   Thanks to Konstantin Baikov in `PR #424 <https://github.com/adamchainz/time-machine/pull/424>`__.
 
+* Fix class decorator for classmethod overrides.
+
+  Thanks to Pavel Bitiukov for the reproducer in `PR #404 <https://github.com/adamchainz/time-machine/pull/404>`__.
+
 * Avoid calling deprecated ``uuid._load_system_functions()`` on Python 3.9+.
 
   Thanks to Nikita Sobolev for the ping in `CPython Issue #113308 <https://github.com/python/cpython/issues/113308>`__.
diff --git a/src/time_machine/__init__.py b/src/time_machine/__init__.py
index ee366f8..877b325 100644
--- a/src/time_machine/__init__.py
+++ b/src/time_machine/__init__.py
@@ -293,24 +293,26 @@ def __call__(
                 raise TypeError("Can only decorate unittest.TestCase subclasses.")
 
             # Modify the setUpClass method
-            orig_setUpClass = wrapped.setUpClass
+            orig_setUpClass = wrapped.setUpClass.__func__  # type: ignore[attr-defined]
 
             @functools.wraps(orig_setUpClass)
             def setUpClass(cls: type[TestCase]) -> None:
                 self.__enter__()
                 try:
-                    orig_setUpClass()
+                    orig_setUpClass(cls)
                 except Exception:
                     self.__exit__(*sys.exc_info())
                     raise
 
             wrapped.setUpClass = classmethod(setUpClass)  # type: ignore[assignment]
 
-            orig_tearDownClass = wrapped.tearDownClass
+            orig_tearDownClass = (
+                wrapped.tearDownClass.__func__  # type: ignore[attr-defined]
+            )
 
             @functools.wraps(orig_tearDownClass)
             def tearDownClass(cls: type[TestCase]) -> None:
-                orig_tearDownClass()
+                orig_tearDownClass(cls)
                 self.__exit__(None, None, None)
 
             wrapped.tearDownClass = classmethod(  # type: ignore[assignment]
diff --git a/tests/test_time_machine.py b/tests/test_time_machine.py
index b014d13..11dc433 100644
--- a/tests/test_time_machine.py
+++ b/tests/test_time_machine.py
@@ -558,6 +558,30 @@ class Something:
     assert excinfo.value.args == ("Can only decorate unittest.TestCase subclasses.",)
 
 
+@time_machine.travel(EPOCH)
+class ClassDecoratorInheritanceBase(TestCase):
+    prop: bool
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.setUpTestData()
+
+    @classmethod
+    def setUpTestData(cls) -> None:
+        cls.prop = True
+
+
+class ClassDecoratorInheritanceTests(ClassDecoratorInheritanceBase):
+    @classmethod
+    def setUpTestData(cls) -> None:
+        super().setUpTestData()
+        cls.prop = False
+
+    def test_ineheritance_correctly_rebound(self):
+        assert self.prop is False
+
+
 class TestMethodDecorator:
     @time_machine.travel(EPOCH + 95.0)
     def test_method_decorator(self):