Avoid retaining objects that are being deallocated.

We do not want to retain objects that are being deallocated because this will cause
a crash later when the retaining object releases them.

An example of this is when a delegate (that is mocked) is called in a dealloc for an
object. The mock retains the deallocating object as part of an invocation and then
things crash later when the invocation is released as part of the mock deallocating
itself.

Author: dmaclach

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -20,7 +20,7 @@
 #import "NSInvocation+OCMAdditions.h"
 #import "OCMFunctionsPrivate.h"
 #import "NSMethodSignature+OCMAdditions.h"
-
+#import "NSObject+OCMAdditions.h"
 
 #if (TARGET_OS_OSX && (!defined(__MAC_10_10) || __MAC_OS_X_VERSION_MIN_REQUIRED < __MAC_10_10)) || \
     (TARGET_OS_IPHONE && (!defined(__IPHONE_8_0) || __IPHONE_OS_VERSION_MIN_REQUIRED < __IPHONE_8_0))
@@ -30,7 +30,6 @@ static BOOL OCMObjectIsClass(id object) {
 #define object_isClass OCMObjectIsClass
 #endif
 
-
 @implementation NSInvocation(OCMAdditions)
 
 + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)arguments
@@ -66,7 +65,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     NSMutableArray *retainedArguments = [[NSMutableArray alloc] init];
 
     id target = [self target];
-    if((target != nil) && (target != objectToExclude) && !object_isClass(target))
+    if((target != nil) && (target != objectToExclude) && !object_isClass(target) && ![target _isDeallocating])
     {
         // Bad things will happen if the target is a block since it's not being
         // copied. There isn't a very good way to tell if an invocation's target
@@ -84,7 +83,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
         {
             id argument;
             [self getArgument:&argument atIndex:index];
-            if((argument != nil) && (argument != objectToExclude))
+            if((argument != nil) && (argument != objectToExclude) && ![argument _isDeallocating])
             {
                 if(OCMIsBlockType(argumentType))
                 {
@@ -106,7 +105,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     {
         id returnValue;
         [self getReturnValue:&returnValue];
-        if((returnValue != nil) && (returnValue != objectToExclude))
+        if((returnValue != nil) && (returnValue != objectToExclude) && ![returnValue _isDeallocating])
         {
             if(OCMIsBlockType(returnType))
             {

Source/OCMock/NSObject+OCMAdditions.h

@@ -22,3 +22,9 @@
 + (void)enumerateMethodsInClass:(Class)aClass usingBlock:(void (^)(Class cls, SEL sel))aBlock;
 
 @end
+
+
+@interface  NSObject (Internals)
+// Return YES if an object is in the process of being deallocated.
+- (BOOL)_isDeallocating;
+@end

Source/OCMock/OCMStubRecorder.m

@@ -24,6 +24,7 @@
 #import "OCMBlockCaller.h"
 #import "OCMRealObjectForwarder.h"
 #import "OCMInvocationStub.h"
+#import "NSObject+OCMAdditions.h"
 
 
 @implementation OCMStubRecorder
@@ -51,7 +52,7 @@ - (OCMInvocationStub *)stub
 - (id)andReturn:(id)anObject
 {
     id action;
-    if(anObject == mockObject)
+    if(anObject == mockObject || [anObject _isDeallocating])
     {
         action = [[[OCMNonRetainingObjectReturnValueProvider alloc] initWithValue:anObject] autorelease];
     }

Source/OCMockTests/OCMockObjectPartialMocksTests.m

@@ -608,5 +608,4 @@ - (void)testMethodSwizzlingWorksForVoidReturns
 	XCTAssertNoThrow([foo method1], @"Should have worked.");
 }
 
-
 @end

Source/OCMockTests/OCMockObjectRuntimeTests.m

@@ -54,14 +54,15 @@ - (NSString *)stringForTypedef:(TypedefString *)string
 
 @interface TestDelegate : NSObject
 
-- (void)go;
+- (id)go:(id)sender;
 
 @end
 
 @implementation TestDelegate
 
-- (void)go
+- (id)go:(id)sender
 {
+    return sender;
 }
 
 @end
@@ -77,7 +78,38 @@ @implementation TestClassWithDelegate
 - (void)run
 {
     TestDelegate *delegate = self.delegate;
-    [delegate go];
+    [delegate go:nil];
+}
+
+@end
+
+
+@interface TestClassThatCallsDelegateOnDealloc : NSObject
+
+@property (nonatomic, weak) TestDelegate *delegate;
+
+@end
+
+@implementation TestClassThatCallsDelegateOnDealloc
+
+- (void)dealloc
+{
+    TestDelegate *delegate = self.delegate;
+    [delegate go:self];
+}
+
+@end
+
+
+@interface TestClassThatCallsMockWithSelfAsReturnInDealloc : NSObject
+@end
+
+@implementation TestClassThatCallsMockWithSelfAsReturnInDealloc
+
+- (void)dealloc
+{
+    id newMock = [OCMockObject mockForClass:[NSString class]];
+    [[[newMock stub] andReturn:self] uppercaseString];
 }
 
 @end
@@ -267,7 +299,7 @@ - (void)testWeakReferencesShouldStayAround
 
     [object run];
 
-    OCMVerify([mockDelegate go]);
+    OCMVerify([mockDelegate go:nil]);
     XCTAssertNotNil(object.delegate, @"Should still have delegate");
 }
 
@@ -283,6 +315,24 @@ - (void)testDynamicSubclassesShouldBeDisposed
     XCTAssertEqual(numClassesBefore, numClassesAfter, @"Should have disposed dynamically generated classes.");
 }
 
+- (void)testHandlesCallingMockWithSelfAsArgumentInDealloc
+{
+    // Note that this test will crash on failure.
+    id mock = [OCMockObject mockForClass:[TestDelegate class]];
+    [[mock expect] go:OCMOCK_ANY];
+    TestClassThatCallsDelegateOnDealloc *foo = [[TestClassThatCallsDelegateOnDealloc alloc] init];
+    foo.delegate = mock;
+    foo = nil;
+    [mock verify];
+}
+
+- (void)testHandlesCallingMockWithSelfAsReturnInDealloc
+{
+    // Note that this test will crash on failure.
+    TestClassThatCallsMockWithSelfAsReturnInDealloc *foo = [[TestClassThatCallsMockWithSelfAsReturnInDealloc alloc] init];
+    foo = nil;
+}
+
 
 #pragma mark    verify mocks work properly when mocking init