Fix for #444.

Author: erikdoe

Source/Changes.txt

@@ -1,6 +1,12 @@
 Listing of notable changes by release. More detail is usually found in the Git
 commit messages and/or the pull requests.
 
+OCMock 3.7.1 (unreleased)
+
+* Fixed a bug that caused double-counting of method invocations on partial
+  mocks under certain circumstances.
+
+
 OCMock 3.7 (2020-07-15)
 
 * Fixed mocking init methods when using ARC. This worked before but could

Source/OCMock/OCMockObject.h

@@ -63,6 +63,7 @@
 
 - (void)addStub:(OCMInvocationStub *)aStub;
 - (void)addExpectation:(OCMInvocationExpectation *)anExpectation;
+- (void)addInvocation:(NSInvocation *)anInvocation;
 
 - (BOOL)handleInvocation:(NSInvocation *)anInvocation;
 - (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation;

Source/OCMock/OCMockObject.m

@@ -134,6 +134,17 @@ - (void)addStub:(OCMInvocationStub *)aStub
     }
 }
 
+- (OCMInvocationStub *)stubForInvocation:(NSInvocation *)anInvocation
+{
+    @synchronized(stubs)
+    {
+        for(OCMInvocationStub *stub in stubs)
+            if([stub matchesInvocation:anInvocation])
+                return stub;
+        return nil;
+    }
+}
+
 - (void)addExpectation:(OCMInvocationExpectation *)anExpectation
 {
     @synchronized(expectations)
@@ -150,6 +161,21 @@ - (void)assertInvocationsArrayIsPresent
     }
 }
 
+- (void)addInvocation:(NSInvocation *)anInvocation
+{
+    @synchronized(invocations)
+    {
+        // We can't do a normal retain arguments on anInvocation because its target/arguments/return
+        // value could be self. That would produce a retain cycle self->invocations->anInvocation->self.
+        // However we need to retain everything on anInvocation that isn't self because we expect them to
+        // stick around after this method returns. Use our special method to retain just what's needed.
+        // This still doesn't completely prevent retain cycles since any of the arguments could have a
+        // strong reference to self. Those will have to be broken with manual calls to -stopMocking.
+        [anInvocation retainObjectArgumentsExcludingObject:self];
+        [invocations addObject:anInvocation];
+    }
+}
+
 
 #pragma mark  Public API
 
@@ -374,37 +400,15 @@ - (void)forwardInvocation:(NSInvocation *)anInvocation
 - (BOOL)handleInvocation:(NSInvocation *)anInvocation
 {
     [self assertInvocationsArrayIsPresent];
-    @synchronized(invocations)
-    {
-        // We can't do a normal retain arguments on anInvocation because its target/arguments/return
-        // value could be self. That would produce a retain cycle self->invocations->anInvocation->self.
-        // However we need to retain everything on anInvocation that isn't self because we expect them to
-        // stick around after this method returns. Use our special method to retain just what's needed.
-        // This still doesn't completely prevent retain cycles since any of the arguments could have a
-        // strong reference to self. Those will have to be broken with manual calls to -stopMocking.
-        [anInvocation retainObjectArgumentsExcludingObject:self];
-        [invocations addObject:anInvocation];
-    }
+    [self addInvocation:anInvocation];
 
-    OCMInvocationStub *stub = nil;
-    @synchronized(stubs)
-    {
-        for(stub in stubs)
-        {
-            // If the stub forwards its invocation to the real object, then we don't want to do
-            // handleInvocation: yet, since forwarding the invocation to the real object could call a
-            // method that is expected to happen after this one, which is bad if expectationOrderMatters
-            // is YES
-            if([stub matchesInvocation:anInvocation])
-                break;
-        }
-        // Retain the stub in case it ends up being removed from stubs and expectations, since we still
-        // have to call handleInvocation on the stub at the end
-        [stub retain];
-    }
+    OCMInvocationStub *stub = [self stubForInvocation:anInvocation];
     if(stub == nil)
         return NO;
 
+    // Retain the stub in case it ends up being removed because we still need it at the end for handleInvocation:
+    [stub retain];
+
     BOOL removeStub = NO;
     @synchronized(expectations)
     {
@@ -418,8 +422,7 @@ - (BOOL)handleInvocation:(NSInvocation *)anInvocation
             }
 
             // We can't check isSatisfied yet, since the stub won't be satisfied until we call
-            // handleInvocation:, and we don't want to call handleInvocation: yes for the reason in the
-            // comment above, since we'll still have the current expectation in the expectations array, which
+            // handleInvocation: since we'll still have the current expectation in the expectations array, which
             // will cause an exception if expectationOrderMatters is YES and we're not ready for any future
             // expected methods to be called yet
             if(![(OCMInvocationExpectation *)stub isMatchAndReject])
@@ -437,9 +440,12 @@ - (BOOL)handleInvocation:(NSInvocation *)anInvocation
         }
     }
 
-    @try {
+    @try
+    {
         [stub handleInvocation:anInvocation];
-    } @finally {
+    }
+    @finally
+    {
         [stub release];
     }
 

Source/OCMock/OCPartialMockObject.h

@@ -18,7 +18,8 @@
 
 @interface OCPartialMockObject : OCClassMockObject 
 {
-	NSObject	*realObject;
+	NSObject		*realObject;
+	NSInvocation	*invocationFromMock;
 }
 
 - (id)initWithObject:(NSObject *)anObject;

Source/OCMock/OCPartialMockObject.m

@@ -105,6 +105,14 @@ - (void)addStub:(OCMInvocationStub *)aStub
         [self setupForwarderForSelector:[[aStub recordedInvocation] selector]];
 }
 
+- (void)addInvocation:(NSInvocation *)anInvocation
+{
+    // If the mock invokes a method on the real object we end up here a second time, but because
+    // the mock has added the invocation already we do not want to add it again.
+    if((invocationFromMock == nil) || ([anInvocation selector] != [invocationFromMock selector]))
+        [super addInvocation:anInvocation];
+}
+
 - (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
 {
 	// In the case of an init that is called on a mock we must return the mock instance and
@@ -118,12 +126,16 @@ - (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
 		targetReceivingInit = [anInvocation target];
 		[realObject retain];
 	}
+
+	invocationFromMock = anInvocation;
 	[anInvocation invokeWithTarget:realObject];
-	if (targetReceivingInit)
+    invocationFromMock = nil;
+
+	if(targetReceivingInit)
 	{
 		id returnVal;
 		[anInvocation getReturnValue:&returnVal];
-		if (returnVal == realObject)
+		if(returnVal == realObject)
 		{
 			[anInvocation setReturnValue:&self];
 			[realObject release];
@@ -247,7 +259,6 @@ - (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation
 }
 
 
-
 #pragma mark    Verification handling
 
 - (NSString *)descriptionForVerificationFailureWithMatcher:(OCMInvocationMatcher *)matcher quantifier:(OCMQuantifier *)quantifier invocationCount:(NSUInteger)count

Source/OCMockTests/OCMockObjectPartialMocksTests.m

@@ -17,6 +17,7 @@
 #import <CoreData/CoreData.h>
 #import <XCTest/XCTest.h>
 #import <OCMock/OCMock.h>
+#import "OCPartialMockObject.h"
 #import "TestClassWithCustomReferenceCounting.h"
 
 #if TARGET_OS_IPHONE
@@ -170,8 +171,6 @@ - (NSString *)categoryMethod
 @end
 
 
-
-
 @interface OCMockObjectPartialMocksTests : XCTestCase
 {
     int numKVOCallbacks;
@@ -209,6 +208,28 @@ @implementation OCTestManagedObject
 @end
 
 
+#pragma mark   Category for testing
+
+@interface OCPartialMockObject(AccessToInvocationsForTesting)
+
+- (NSArray *)invocationsExcludingInitialize;
+
+@end
+
+@implementation OCPartialMockObject(AccessToInvocationsForTesting)
+
+- (NSArray *)invocationsExcludingInitialize
+{
+	NSMutableArray *filteredInvocations = [[NSMutableArray alloc] init];
+	for(NSInvocation *i in invocations)
+		if([NSStringFromSelector([i selector]) hasSuffix:@"initialize"] == NO)
+			[filteredInvocations addObject:i];
+
+	return filteredInvocations;
+}
+
+@end
+
 
 @implementation OCMockObjectPartialMocksTests
 
@@ -280,6 +301,64 @@ - (void)testInvocationsOfNSObjectCategoryMethodsCanBeStubbed
 }
 
 
+#pragma mark   Tests for remembering invocations for later verification
+
+- (void)testRecordsInvocationWhenRealObjectIsUsed
+{
+	TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
+	id mock = [OCMockObject partialMockForObject:realObject];
+
+	[realObject foo];
+
+	XCTAssertEqual(1, [[mock invocationsExcludingInitialize] count]);
+}
+
+- (void)testRecordsInvocationWhenMockIsUsed
+{
+	TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
+	id mock = [OCMockObject partialMockForObject:realObject];
+
+	[mock foo];
+
+	XCTAssertEqual(1, [[mock invocationsExcludingInitialize] count]);
+}
+
+- (void)testRecordsInvocationWhenRealObjectIsUsedAndMethodIsStubbed
+{
+	TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
+	id mock = [OCMockObject partialMockForObject:realObject];
+	[[[mock stub] andReturn:@"bar"] foo];
+
+	id res = [realObject foo];
+
+	XCTAssertEqualObjects(@"bar", res);
+	XCTAssertEqual(1, [[mock invocationsExcludingInitialize] count]);
+}
+
+- (void)testRecordsInvocationWhenMockIsUsedAndMethodIsStubbed
+{
+	TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
+	id mock = [OCMockObject partialMockForObject:realObject];
+	[[[mock stub] andReturn:@"bar"] foo];
+
+	id res = [mock foo];
+
+	XCTAssertEqualObjects(@"bar", res);
+	XCTAssertEqual(1, [[mock invocationsExcludingInitialize] count]);
+}
+
+- (void)testRecordsInvocationWhenMockIsUsedAndMethodIsStubbedAndForwardsToRealObject
+{
+	TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
+	id mock = [OCMockObject partialMockForObject:realObject];
+	[[[mock stub] andForwardToRealObject] foo];
+
+	[mock foo];
+
+	XCTAssertEqual(1, [[mock invocationsExcludingInitialize] count]);
+}
+
+
 #pragma mark   Tests for behaviour when setting up partial mocks
 
 - (void)testPartialMockClassOverrideReportsOriginalClass