Update OCMock to 3.1.5

third_party/ocmock/OCMock/OCPartialMockObject.m

Index: third_party/ocmock/OCMock/OCPartialMockObject.m
diff --git a/third_party/ocmock/OCMock/OCPartialMockObject.m b/third_party/ocmock/OCMock/OCPartialMockObject.m
index fb0b8bd592bc6298b58e8a7491860d509d622c41..82f376ab4954697b67da1e0a5c0ca22e87f0102f 100644
--- a/third_party/ocmock/OCMock/OCPartialMockObject.m
+++ b/third_party/ocmock/OCMock/OCPartialMockObject.m
@@ -1,74 +1,52 @@
-//---------------------------------------------------------------------------------------
-//  $Id$
-//  Copyright (c) 2009 by Mulle Kybernetik. See License file for details.
-//---------------------------------------------------------------------------------------
+/*
+ *  Copyright (c) 2009-2015 Erik Doernenburg and contributors
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use these files except in compliance with the License. You may obtain
+ *  a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  License for the specific language governing permissions and limitations
+ *  under the License.
+ */
 
 #import <objc/runtime.h>
-#import "OCPartialMockRecorder.h"
+#import "OCMockObject.h"
 #import "OCPartialMockObject.h"
+#import "NSMethodSignature+OCMAdditions.h"
+#import "NSObject+OCMAdditions.h"
+#import "OCMFunctions.h"
+#import "OCMInvocationStub.h"
 
 
-@interface OCPartialMockObject (Private)
-- (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation;
-@end 
-
-
-NSString *OCMRealMethodAliasPrefix = @"ocmock_replaced_";
-
 @implementation OCPartialMockObject
 
-
-#pragma mark  Mock table
-
-static NSMutableDictionary *mockTable;
-
-+ (void)initialize
-{
-	if(self == [OCPartialMockObject class])
-		mockTable = [[NSMutableDictionary alloc] init];
-}
-
-+ (void)rememberPartialMock:(OCPartialMockObject *)mock forObject:(id)anObject
-{
-	[mockTable setObject:[NSValue valueWithNonretainedObject:mock] forKey:[NSValue valueWithNonretainedObject:anObject]];
-}
-
-+ (void)forgetPartialMockForObject:(id)anObject
-{
-	[mockTable removeObjectForKey:[NSValue valueWithNonretainedObject:anObject]];
-}
-
-+ (OCPartialMockObject *)existingPartialMockForObject:(id)anObject
-{
-	OCPartialMockObject *mock = [[mockTable objectForKey:[NSValue valueWithNonretainedObject:anObject]] nonretainedObjectValue];
-	if(mock == nil)
-		[NSException raise:NSInternalInconsistencyException format:@"No partial mock for object %p", anObject];
-	return mock;
-}
-
-
-
 #pragma mark  Initialisers, description, accessors, etc.
 
 - (id)initWithObject:(NSObject *)anObject
 {
+    NSParameterAssert(anObject != nil);
+    [self assertClassIsSupported:[anObject class]];
 	[super initWithClass:[anObject class]];
 	realObject = [anObject retain];
-	[[self class] rememberPartialMock:self forObject:anObject];
-	[self setupSubclassForObject:realObject];
+    [self prepareObjectForInstanceMethodMocking];
 	return self;
 }
 
 - (void)dealloc
 {
-	if(realObject != nil)
-		[self stop];
+	[self stopMocking];
+	[realObject release];
 	[super dealloc];
 }
 
 - (NSString *)description
 {
-	return [NSString stringWithFormat:@"OCPartialMockObject[%@]", NSStringFromClass(mockedClass)];
+	return [NSString stringWithFormat:@"OCPartialMockObject(%@)", NSStringFromClass(mockedClass)];
 }
 
 - (NSObject *)realObject
@@ -76,66 +54,164 @@ static NSMutableDictionary *mockTable;
 	return realObject;
 }
 
-- (void)stop
+#pragma mark  Helper methods
+
+- (void)assertClassIsSupported:(Class)class
 {
-	object_setClass(realObject, [self mockedClass]);
-	[realObject release];
-	[[self class] forgetPartialMockForObject:realObject];
-	realObject = nil;
+    NSString *classname = NSStringFromClass(class);
+    NSString *reason = nil;
+    if([classname hasPrefix:@"__NSTagged"] || [classname hasPrefix:@"NSTagged"])
+        reason = [NSString stringWithFormat:@"OCMock does not support partially mocking tagged classes; got %@", classname];
+    else if([classname hasPrefix:@"__NSCF"])
+        reason = [NSString stringWithFormat:@"OCMock does not support partially mocking toll-free bridged classes; got %@", classname];
+
+    if(reason != nil)
+        [[NSException exceptionWithName:NSInvalidArgumentException reason:reason userInfo:nil] raise];
 }
 
 
-#pragma mark  Subclass management
+#pragma mark  Extending/overriding superclass behaviour
+
+- (void)stopMocking
+{
+    if(realObject != nil)
+    {
+        OCMSetAssociatedMockForObject(nil, realObject);
+        object_setClass(realObject, [self mockedClass]);
+        [realObject release];
+        realObject = nil;
+    }
+    [super stopMocking];
+}
 
-- (void)setupSubclassForObject:(id)anObject
+- (void)addStub:(OCMInvocationStub *)aStub
 {
-	Class realClass = [anObject class];
-	double timestamp = [NSDate timeIntervalSinceReferenceDate];
-	const char *className = [[NSString stringWithFormat:@"%@-%p-%f", realClass, anObject, timestamp] UTF8String];
-	Class subclass = objc_allocateClassPair(realClass, className, 0);
-	objc_registerClassPair(subclass);
-	object_setClass(anObject, subclass);
-	
-	Method forwardInvocationMethod = class_getInstanceMethod([self class], @selector(forwardInvocationForRealObject:));
-	IMP forwardInvocationImp = method_getImplementation(forwardInvocationMethod);
-	const char *forwardInvocationTypes = method_getTypeEncoding(forwardInvocationMethod);
-	class_addMethod(subclass, @selector(forwardInvocation:), forwardInvocationImp, forwardInvocationTypes);
+    [super addStub:aStub];
+    if(![aStub recordedAsClassMethod])
+        [self setupForwarderForSelector:[[aStub recordedInvocation] selector]];
 }
 
-- (void)setupForwarderForSelector:(SEL)selector
+- (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
 {
-	Class subclass = [[self realObject] class];
-	Method originalMethod = class_getInstanceMethod([subclass superclass], selector);
-	IMP originalImp = method_getImplementation(originalMethod);
+	[anInvocation invokeWithTarget:realObject];
+}
+
 
-	IMP forwarderImp = [subclass instanceMethodForSelector:@selector(aMethodThatMustNotExist)];
-	class_addMethod(subclass, method_getName(originalMethod), forwarderImp, method_getTypeEncoding(originalMethod)); 
+#pragma mark  Subclass management
 
-	SEL aliasSelector = NSSelectorFromString([OCMRealMethodAliasPrefix stringByAppendingString:NSStringFromSelector(selector)]);
-	class_addMethod(subclass, aliasSelector, originalImp, method_getTypeEncoding(originalMethod));
+- (void)prepareObjectForInstanceMethodMocking
+{
+    OCMSetAssociatedMockForObject(self, realObject);
+
+    /* dynamically create a subclass and set it as the class of the object */
+    Class subclass = OCMCreateSubclass(mockedClass, realObject);
+	object_setClass(realObject, subclass);
+
+    /* point forwardInvocation: of the object to the implementation in the mock */
+	Method myForwardMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardInvocationForRealObject:));
+	IMP myForwardIMP = method_getImplementation(myForwardMethod);
+    class_addMethod(subclass, @selector(forwardInvocation:), myForwardIMP, method_getTypeEncoding(myForwardMethod));
+
+    /* do the same for forwardingTargetForSelector, remember existing imp with alias selector */
+    Method myForwardingTargetMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardingTargetForSelectorForRealObject:));
+    IMP myForwardingTargetIMP = method_getImplementation(myForwardingTargetMethod);
+    IMP originalForwardingTargetIMP = [mockedClass instanceMethodForSelector:@selector(forwardingTargetForSelector:)];
+    class_addMethod(subclass, @selector(forwardingTargetForSelector:), myForwardingTargetIMP, method_getTypeEncoding(myForwardingTargetMethod));
+    class_addMethod(subclass, @selector(ocmock_replaced_forwardingTargetForSelector:), originalForwardingTargetIMP, method_getTypeEncoding(myForwardingTargetMethod));
+
+    /* We also override the -class method to return the original class */
+    Method myObjectClassMethod = class_getInstanceMethod([self mockObjectClass], @selector(classForRealObject));
+    const char *objectClassTypes = method_getTypeEncoding(myObjectClassMethod);
+    IMP myObjectClassImp = method_getImplementation(myObjectClassMethod);
+    class_addMethod(subclass, @selector(class), myObjectClassImp, objectClassTypes);
+
+    /* Adding forwarder for most instance methods to allow for verify after run */
+    NSArray *methodBlackList = @[@"class", @"forwardingTargetForSelector:", @"methodSignatureForSelector:", @"forwardInvocation:",
+            @"allowsWeakReference", @"retainWeakReference", @"isBlock"];
+    [NSObject enumerateMethodsInClass:mockedClass usingBlock:^(Class cls, SEL sel) {
+        if((cls == [NSObject class]) || (cls == [NSProxy class]))
+            return;
+        NSString *className = NSStringFromClass(cls);
+        NSString *selName = NSStringFromSelector(sel);
+        if(([className hasPrefix:@"NS"] || [className hasPrefix:@"UI"]) &&
+           ([selName hasPrefix:@"_"] || [selName hasSuffix:@"_"]))
+            return;
+        if([methodBlackList containsObject:selName])
+            return;
+        @try
+        {
+            [self setupForwarderForSelector:sel];
+        }
+        @catch(NSException *e)
+        {
+            // ignore for now
+        }
+    }];
 }
 
-- (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation
+- (void)setupForwarderForSelector:(SEL)sel
 {
-	// in here "self" is a reference to the real object, not the mock
-	OCPartialMockObject *mock = [OCPartialMockObject existingPartialMockForObject:self];
-	if([mock handleInvocation:anInvocation] == NO)
-		[NSException raise:NSInternalInconsistencyException format:@"Ended up in subclass forwarder for %@ with unstubbed method %@",
-		 [self class], NSStringFromSelector([anInvocation selector])];
+    SEL aliasSelector = OCMAliasForOriginalSelector(sel);
+    if(class_getInstanceMethod(object_getClass(realObject), aliasSelector) != NULL)
+        return;
+
+    Method originalMethod = class_getInstanceMethod(mockedClass, sel);
+	IMP originalIMP = method_getImplementation(originalMethod);
+    const char *types = method_getTypeEncoding(originalMethod);
+    /* Might be NULL if the selector is forwarded to another class */
+    // TODO: check the fallback implementation is actually sufficient
+    if(types == NULL)
+        types = ([[mockedClass instanceMethodSignatureForSelector:sel] fullObjCTypes]);
+
+    Class subclass = object_getClass([self realObject]);
+    IMP forwarderIMP = [mockedClass instanceMethodForwarderForSelector:sel];
+    class_replaceMethod(subclass, sel, forwarderIMP, types);
+	class_addMethod(subclass, aliasSelector, originalIMP, types);
 }
 
 
+// Implementation of the -class method; return the Class that was reported with [realObject class] prior to mocking
+- (Class)classForRealObject
+{
+    // in here "self" is a reference to the real object, not the mock
+    OCPartialMockObject *mock = OCMGetAssociatedMockForObject(self);
+    if(mock == nil)
+        [NSException raise:NSInternalInconsistencyException format:@"No partial mock for object %p", self];
+    return [mock mockedClass];
+}
 
-#pragma mark  Overrides
 
-- (id)getNewRecorder
+- (id)forwardingTargetForSelectorForRealObject:(SEL)sel
 {
-	return [[[OCPartialMockRecorder alloc] initWithSignatureResolver:self] autorelease];
+	// in here "self" is a reference to the real object, not the mock
+    OCPartialMockObject *mock = OCMGetAssociatedMockForObject(self);
+    if(mock == nil)
+        [NSException raise:NSInternalInconsistencyException format:@"No partial mock for object %p", self];
+    if([mock handleSelector:sel])
+        return self;
+
+    return [self ocmock_replaced_forwardingTargetForSelector:sel];
 }
 
-- (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
+//  Make the compiler happy in -forwardingTargetForSelectorForRealObject: because it can't find the message…
+- (id)ocmock_replaced_forwardingTargetForSelector:(SEL)sel
 {
-	[anInvocation invokeWithTarget:realObject];
+    return nil;
+}
+
+
+- (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation
+{
+	// in here "self" is a reference to the real object, not the mock
+    OCPartialMockObject *mock = OCMGetAssociatedMockForObject(self);
+    if(mock == nil)
+        [NSException raise:NSInternalInconsistencyException format:@"No partial mock for object %p", self];
+
+	if([mock handleInvocation:anInvocation] == NO)
+    {
+        [anInvocation setSelector:OCMAliasForOriginalSelector([anInvocation selector])];
+        [anInvocation invoke];
+    }
 }