Fix for Issue 5, submitted by agoratim.

1 files changed, 5 insertions, 2 deletions

diff --git a/stubout.py b/stubout.py
index 6d04996..a452241 100644
--- a/stubout.py
+++ b/stubout.py
@@ -118,8 +118,11 @@ class StubOutForTesting:
     old_child = getattr(parent, child_name)
 
     old_attribute = parent.__dict__.get(child_name)
-    if old_attribute is not None and isinstance(old_attribute, staticmethod):
-      old_child = staticmethod(old_child)
+    if old_attribute is not None:
+      if isinstance(old_attribute, staticmethod):
+        old_child = staticmethod(old_child)
+      elif isinstance(old_attribute, classmethod):
+        old_child = classmethod(old_child.im_func)
 
     self.cache.append((parent, old_child, child_name))
     setattr(parent, child_name, new_child)