Automatically stop all live mocks at the end of each test case/suite

If the user is using XCTest with OCMock, this registers a test observer that takes care
of stopping all live mocks appropriately.

For mocks that are created in +setUp, those will get stopped at the end of the suite.
For mocks that are created in -setUp or in test cases themselves, those will get
stopped at the end of the testcase.

While these mocks are being stopped and testcases/suites are being torndown, messages
sent to mocks are not going to trigger the exception about calling a mock after it has
had stopMocking called on it. This allows objects that may refer to mocks in dealloc
methods to be cleaned up in autoreleasepools or due to stopMocking being called
without the mocks throwing exceptions.

This should greatly simplify cleaning up mocks and remove a lot of potential leakage.
It also makes sure that class mocks that mock class methods will not persist across tests.

Author: dmaclach

Source/OCMock.xcodeproj/project.pbxproj

@@ -294,6 +294,8 @@
 		8BF7402124772B0600B9A52C /* XCTest.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 03565A1D18F05626003AE91E /* XCTest.framework */; };
 		8BF7402224772B0800B9A52C /* XCTest.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 03565A1D18F05626003AE91E /* XCTest.framework */; };
 		8BF7402324772B0800B9A52C /* XCTest.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 03565A1D18F05626003AE91E /* XCTest.framework */; };
+		8BC0A67C242D08D800695F71 /* OCMockObjectCleanupTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BC0A67B242D08D800695F71 /* OCMockObjectCleanupTests.m */; };
+		8BC0A67D242D08E400695F71 /* OCMockObjectCleanupTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 8BC0A67B242D08D800695F71 /* OCMockObjectCleanupTests.m */; };
 		8DE97C5522B43EE60098C63F /* OCMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3159E146333BF0052CD09 /* OCMockObject.m */; };
 		8DE97C5622B43EE60098C63F /* OCClassMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3158C146333BF0052CD09 /* OCClassMockObject.m */; };
 		8DE97C5722B43EE60098C63F /* OCPartialMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B315AA146333BF0052CD09 /* OCPartialMockObject.m */; };
@@ -589,6 +591,7 @@
 		8B11D4B92448E53600247BE2 /* OCMCPlusPlus11Tests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = OCMCPlusPlus11Tests.mm; sourceTree = "<group>"; };
 		8B3786A724E5BD5600FD1B5B /* OCMFunctionsTests.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = OCMFunctionsTests.m; sourceTree = "<group>"; };
 		8BF73E52246CA75E00B9A52C /* OCMNoEscapeBlockTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMNoEscapeBlockTests.m; sourceTree = "<group>"; };
+		8BC0A67B242D08D800695F71 /* OCMockObjectCleanupTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMockObjectCleanupTests.m; sourceTree = "<group>"; };
 		8DE97CA022B43EE60098C63F /* OCMock.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = OCMock.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		A02926811CA0725A00594AAF /* TestObjects.xcdatamodel */ = {isa = PBXFileReference; lastKnownFileType = wrapper.xcdatamodel; path = TestObjects.xcdatamodel; sourceTree = "<group>"; };
 		D31108AD1828DB8700737925 /* OCMockLibTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = OCMockLibTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
@@ -755,6 +758,7 @@
 				03AC5C1416DF9FA500D82ECD /* OCMockObjectPartialMocksTests.m */,
 				039F91C516EFB493006C3D70 /* OCMockObjectClassMethodMockingTests.m */,
 				2FA286BFBD8B9D068B41E7EF /* OCMockObjectProtocolMocksTests.m */,
+				8BC0A67B242D08D800695F71 /* OCMockObjectCleanupTests.m */,
 				2FA28EE3142412BD601026EF /* OCMockObjectDynamicPropertyMockingTests.m */,
 				03E98D4F18F310EE00522D42 /* OCMockObjectMacroTests.m */,
 				0354D71F16F23AF5001766BB /* OCMockObjectForwardingTargetTests.m */,
@@ -1536,6 +1540,7 @@
 				03565A4218F05721003AE91E /* OCMockObjectPartialMocksTests.m in Sources */,
 				03565A4C18F05721003AE91E /* NSMethodSignatureOCMAdditionsTests.m in Sources */,
 				03565A4818F05721003AE91E /* OCMStubRecorderTests.m in Sources */,
+				8BC0A67C242D08D800695F71 /* OCMockObjectCleanupTests.m in Sources */,
 				03565A4518F05721003AE91E /* OCMockObjectForwardingTargetTests.m in Sources */,
 				2FA28FA53C57236B6DD64E82 /* OCMockObjectRuntimeTests.m in Sources */,
 				8B11D4BA2448E53600247BE2 /* OCMCPlusPlus11Tests.mm in Sources */,
@@ -1652,6 +1657,7 @@
 				D31108CA1828DBD600737925 /* NSInvocationOCMAdditionsTests.m in Sources */,
 				03C9CA1F18F05A8E006DF94D /* NSMethodSignatureOCMAdditionsTests.m in Sources */,
 				03C9CA1D18F05A75006DF94D /* OCMockObjectProtocolMocksTests.m in Sources */,
+				8BC0A67D242D08E400695F71 /* OCMockObjectCleanupTests.m in Sources */,
 				03E98D5118F310EE00522D42 /* OCMockObjectMacroTests.m in Sources */,
 				A06930951CA1BFC900513023 /* TestObjects.xcdatamodeld in Sources */,
 				8B11D4BB2448E53600247BE2 /* OCMCPlusPlus11Tests.mm in Sources */,

Source/OCMock/OCClassMockObject.m

@@ -32,10 +32,18 @@ @implementation OCClassMockObject
 
 - (id)initWithClass:(Class)aClass
 {
-	[self assertClassIsSupported:aClass];
-	[super init];
-	mockedClass = aClass;
+  @try
+  {
+    [self assertClassIsSupported:aClass];
+    [super init];
+    mockedClass = aClass;
     [self prepareClassForClassMethodMocking];
+  }
+  @catch(NSException *e)
+  {
+      [OCMockObject removeAMockToStop:self];
+      [e raise];
+  }
 	return self;
 }
 

Source/OCMock/OCMInvocationExpectation.m

@@ -16,7 +16,7 @@
 
 #import "OCMInvocationExpectation.h"
 #import "NSInvocation+OCMAdditions.h"
-
+#import "OCMockObject.h"
 
 @implementation OCMInvocationExpectation
 
@@ -52,7 +52,7 @@ - (void)handleInvocation:(NSInvocation *)anInvocation
     if(matchAndReject)
     {
         isSatisfied = NO;
-        [NSException raise:NSInternalInconsistencyException format:@"%@: explicitly disallowed method invoked: %@",
+        [OCMockObject logMatcherIssue:@"%@: explicitly disallowed method invoked: %@",
                 [self description], [anInvocation invocationDescription]];
     }
     else

Source/OCMock/OCMockObject.h

@@ -74,5 +74,8 @@
 - (void)verifyInvocation:(OCMInvocationMatcher *)matcher withQuantifier:(OCMQuantifier *)quantifier atLocation:(OCMLocation *)location;
 - (NSString *)descriptionForVerificationFailureWithMatcher:(OCMInvocationMatcher *)matcher quantifier:(OCMQuantifier *)quantifier invocationCount:(NSUInteger)count;
 
++ (void)removeAMockToStop:(OCMockObject *)mock;
+
++ (void)logMatcherIssue:(NSString *)format, ... NS_FORMAT_FUNCTION(1,2);
 @end
 

Source/OCMock/OCMockObject.m

@@ -29,18 +29,73 @@
 #import "OCMFunctionsPrivate.h"
 #import "NSInvocation+OCMAdditions.h"
 
+@class XCTestCase;
+@class XCTest;
+
+// gMocksToStopRecorders is a stack of recorders that gets added to and removed from
+// as we enter test suite/case scopes.
+// Controlled by OCMockXCTestObserver.
+static NSMutableArray<NSHashTable<OCMockObject *> *> *gMocksToStopRecorders;
+
+// Flag that controls whether we should be asserting after stopmocking is called.
+// Controlled by OCMockXCTestObserver.
+static BOOL gAssertOnCallsAfterStopMocking;
+
+// Flag that tracks if we are stopping the mocks.
+static BOOL gStoppingMocks = NO;
 
 @implementation OCMockObject
 
 #pragma mark  Class initialisation
 
 + (void)initialize
 {
-	if([[NSInvocation class] instanceMethodSignatureForSelector:@selector(getArgumentAtIndexAsObject:)] == NULL)
-		[NSException raise:NSInternalInconsistencyException format:@"** Expected method not present; the method getArgumentAtIndexAsObject: is not implemented by NSInvocation. If you see this exception it is likely that you are using the static library version of OCMock and your project is not configured correctly to load categories from static libraries. Did you forget to add the -ObjC linker flag?"];
+    if([[NSInvocation class] instanceMethodSignatureForSelector:@selector(getArgumentAtIndexAsObject:)] == NULL)
+    {
+        [NSException raise:NSInternalInconsistencyException format:@"** Expected method not present; the method getArgumentAtIndexAsObject: is not implemented by NSInvocation. If you see this exception it is likely that you are using the static library version of OCMock and your project is not configured correctly to load categories from static libraries. Did you forget to add the -ObjC linker flag?"];
+    }
+}
+
+#pragma mark  Mock cleanup recording
+
++ (void)recordAMockToStop:(OCMockObject *)mock
+{
+  @synchronized(self)
+  {
+    if(gStoppingMocks)
+    {
+      [NSException raise:NSInternalInconsistencyException format:@"Attempting to add a mock while mocks are being stopped."];
+    }
+    [[gMocksToStopRecorders lastObject] addObject:mock];
+  }
 }
 
++ (void)removeAMockToStop:(OCMockObject *)mock
+{
+  @synchronized(self)
+  {
+    if(gStoppingMocks)
+    {
+      [NSException raise:NSInternalInconsistencyException format:@"Attempting to remove a mock while mocks are being stopped."];
+    }
+    [[gMocksToStopRecorders lastObject] removeObject:mock];
+  }
+}
 
++ (void)stopAllCurrentMocks
+{
+  @synchronized(self)
+  {
+    gStoppingMocks = YES;
+    NSHashTable<OCMockObject *> *recorder = [gMocksToStopRecorders lastObject];
+    for (OCMockObject *mock in recorder)
+    {
+      [mock stopMocking];
+    }
+    [recorder removeAllObjects];
+    gStoppingMocks = NO;
+  }
+}
 #pragma mark  Factory methods
 
 + (id)mockForClass:(Class)aClass
@@ -112,6 +167,7 @@ - (instancetype)init
 	expectations = [[NSMutableArray alloc] init];
 	exceptions = [[NSMutableArray alloc] init];
     invocations = [[NSMutableArray alloc] init];
+    [OCMockObject recordAMockToStop:self];
     return self;
 }
 
@@ -161,7 +217,7 @@ - (void)assertInvocationsArrayIsPresent
 {
     if(invocations == nil)
     {
-        [NSException raise:NSInternalInconsistencyException format:@"** Cannot use mock object %@ at %p. This error usually occurs when a mock object is used after stopMocking has been called on it. In most cases it is not necessary to call stopMocking. If you know you have to, please make sure that the mock object is not used afterwards.", [self description], (void *)self];
+        [OCMockObject logMatcherIssue:@"** Cannot use mock object %@ at %p. This error usually occurs when a mock object is used after stopMocking has been called on it. In most cases it is not necessary to call stopMocking. If you know you have to, please make sure that the mock object is not used afterwards.", [self description], (void *)self];
     }
 }
 
@@ -180,6 +236,16 @@ - (void)addInvocation:(NSInvocation *)anInvocation
     }
 }
 
++ (void)logMatcherIssue:(NSString *)format, ...
+{
+    if(gAssertOnCallsAfterStopMocking)
+    {
+        va_list args;
+        va_start(args, format);
+        [NSException raise:NSInternalInconsistencyException format:format arguments:args];
+        va_end(args);
+    }
+}
 
 #pragma mark  Public API
 
@@ -469,7 +535,7 @@ - (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
 {
 	if(isNice == NO)
 	{
-		[NSException raise:NSInternalInconsistencyException format:@"%@: unexpected method invoked: %@ %@",
+		  [OCMockObject logMatcherIssue:@"%@: unexpected method invoked: %@ %@",
                         [self description], [anInvocation invocationDescription], [self _stubDescriptions:NO]];
 	}
 }
@@ -529,4 +595,83 @@ - (NSString *)_stubDescriptions:(BOOL)onlyExpectations
 }
 
 
+@end
+
+/**
+ * The observer gets installed the first time a mock object is created (see +[OCMockObject initialize]
+ * It stops all the mocks that are still active when the testcase has finished.
+ * In many cases this should break a lot of retain loops and allow mocks to be freed.
+ * More importantly this will remove mocks that have mocked a class method and persist across testcases.
+ * It intentionally turns off the assert that fires when calling a mock after stopMocking has been
+ * called on it, because when we are doing cleanup there are cases in dealloc methods where a mock
+ * may be called. We allow the "assert off" state to persist beyond the end of -testCaseDidFinish
+ * because objects may be destroyed by the autoreleasepool that wraps the entire test and this may
+ * cause  mocks to be called. The state is global (instead of per mock) because we want to be able
+ * to catch the case where a mock is trapped by some global state (e.g. a non-mock singleton) and
+ * then that singleton is used in a later test and attempts to call a stopped mock.
+ **/
+@interface OCMockXCTestObserver : NSObject
+@end
+
+// "Fake" Protocol so we can avoid having to link to XCTest, but not get warnings about
+// methods not being declared.
+@protocol OCMockXCTestObservation
++ (id)sharedTestObservationCenter;
+- (void)addTestObserver:(id)observer;
+@end
+
+@implementation OCMockXCTestObserver
+
++ (void)load
+{
+    gMocksToStopRecorders = [[NSMutableArray alloc] init];
+    gAssertOnCallsAfterStopMocking = YES;
+    Class xctest = NSClassFromString(@"XCTestObservationCenter");
+    if (xctest)
+    {
+        // If XCTest is available, we set up an observer to stop our mocks for us.
+        [[xctest sharedTestObservationCenter] addTestObserver:[[OCMockXCTestObserver alloc] init]];
+    }
+}
+
+- (BOOL)conformsToProtocol:(Protocol *)aProtocol
+{
+    // This allows us to avoid linking XCTest into OCMock.
+    return strcmp(protocol_getName(aProtocol), "XCTestObservation") == 0;
+}
+
+- (void)addRecorder
+{
+    gAssertOnCallsAfterStopMocking = YES;
+    NSHashTable<OCMockObject *> *recorder = [NSHashTable weakObjectsHashTable];
+    [gMocksToStopRecorders addObject:recorder];
+}
+
+- (void)finalizeRecorder
+{
+    gAssertOnCallsAfterStopMocking = NO;
+    [OCMockObject stopAllCurrentMocks];
+    [gMocksToStopRecorders removeLastObject];
+}
+
+- (void)testSuiteWillStart:(XCTestCase *)testCase
+{
+    [self addRecorder];
+}
+
+- (void)testSuiteDidFinish:(XCTestCase *)testCase
+{
+    [self finalizeRecorder];
+}
+
+- (void)testCaseWillStart:(XCTestCase *)testCase
+{
+    [self addRecorder];
+}
+
+- (void)testCaseDidFinish:(XCTestCase *)testCase
+{
+    [self finalizeRecorder];
+}
+
 @end

Source/OCMock/OCPartialMockObject.m

@@ -29,12 +29,20 @@ @implementation OCPartialMockObject
 
 - (id)initWithObject:(NSObject *)anObject
 {
+  @try
+  {
      if(anObject == nil)
         [NSException raise:NSInvalidArgumentException format:@"Object cannot be nil."];
-    Class const class = [self classToSubclassForObject:anObject];
-	[super initWithClass:class];
-	realObject = [anObject retain];
+     Class const class = [self classToSubclassForObject:anObject];
+	  [super initWithClass:class];
+	  realObject = [anObject retain];
     [self prepareObjectForInstanceMethodMocking];
+  }
+  @catch(NSException *e)
+  {
+      [OCMockObject removeAMockToStop:self];
+      [e raise];
+  }
 	return self;
 }
 

Source/OCMock/OCProtocolMockObject.m

@@ -24,12 +24,19 @@ @implementation OCProtocolMockObject
 
 - (id)initWithProtocol:(Protocol *)aProtocol
 {
-	if(aProtocol == nil)
-		[NSException raise:NSInvalidArgumentException format:@"Protocol cannot be nil."];
-
-	[super init];
-	mockedProtocol = aProtocol;
-	return self;
+    @try
+    {
+        if(aProtocol == nil)
+          [NSException raise:NSInvalidArgumentException format:@"Protocol cannot be nil."];
+        [super init];
+        mockedProtocol = aProtocol;
+    }
+    @catch(NSException *e)
+    {
+        [OCMockObject removeAMockToStop:self];
+        [e raise];
+    }
+    return self;
 }
 
 - (NSString *)description