Fix multiple expectations for the same classmethod are not matched

src/flexmock/api.py

@@ -244,19 +244,22 @@ def _create_placeholder_mock_for_proper_teardown(
     def _update_method(self, expectation: "Expectation", name: str) -> None:
         method_instance = self._create_mock_method(name)
         obj = self._object
-        if self._hasattr(obj, name) and not hasattr(expectation, "_original"):
-            expectation._update_original(name, obj)
-            method_type = type(_getattr(expectation, "_original"))
-            try:
-                # TODO(herman): this is awful, fix this properly.
-                # When a class/static method is mocked out on an *instance*
-                # we need to fetch the type from the class
-                method_type = type(_getattr(obj.__class__, name))
-            except Exception:
-                pass
-            if method_type in SPECIAL_METHODS:
-                expectation._original_function = getattr(obj, name)
-            expectation._method_type = method_type
+        if self._hasattr(obj, name):
+            if hasattr(expectation, "_original"):
+                expectation._method_type = type(_getattr(expectation, "_original"))
+            else:
+                expectation._update_original(name, obj)
+                method_type = type(_getattr(expectation, "_original"))
+                try:
+                    # TODO(herman): this is awful, fix this properly.
+                    # When a class/static method is mocked out on an *instance*
+                    # we need to fetch the type from the class
+                    method_type = type(_getattr(obj.__class__, name))
+                except Exception:
+                    pass
+                if method_type in SPECIAL_METHODS:
+                    expectation._original_function = getattr(obj, name)
+                expectation._method_type = method_type
         if not inspect.isclass(obj) or expectation._method_type in SPECIAL_METHODS:
             method_instance = types.MethodType(method_instance, obj)
         override = _setattr(obj, name, method_instance)

tests/flexmock_test.py

@@ -562,6 +562,25 @@ def bar(self):
         flexmock(User).should_call("bar").once()
         assert_equal("value", user1.bar())
 
+    def test_with_args_with_instance_method(self):
+        flexmock(SomeClass).should_receive("instance_method").with_args("red").once()
+        flexmock(SomeClass).should_receive("instance_method").with_args("blue").once()
+        instance = SomeClass()
+        instance.instance_method("red")
+        instance.instance_method("blue")
+
+    def test_with_args_with_class_method(self):
+        flexmock(SomeClass).should_receive("class_method").with_args("red").once()
+        flexmock(SomeClass).should_receive("class_method").with_args("blue").once()
+        SomeClass.class_method("red")
+        SomeClass.class_method("blue")
+
+    def test_with_args_with_static_method(self):
+        flexmock(SomeClass).should_receive("static_method").with_args("red").once()
+        flexmock(SomeClass).should_receive("static_method").with_args("blue").once()
+        SomeClass.static_method("red")
+        SomeClass.static_method("blue")
+
     def test_flexmock_should_not_blow_up_on_should_call_for_class_methods(self):
         class User:
             @classmethod