Update OCMock to 3.1.5

third_party/ocmock/OCMock/OCPartialMockObject.m

-//---------------------------------------------------------------------------------------
-//  $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
 {
         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)setupSubclassForObject:(id)anObject
+- (void)stopMocking
 {
-        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);
+    if(realObject != nil)
+    {
+        OCMSetAssociatedMockForObject(nil, realObject);
+        object_setClass(realObject, [self mockedClass]);
+        [realObject release];
+        realObject = nil;
+    }
+    [super stopMocking];
 }
 
-- (void)setupForwarderForSelector:(SEL)selector
+- (void)addStub:(OCMInvocationStub *)aStub
 {
-        Class subclass = [[self realObject] class];
-        Method originalMethod = class_getInstanceMethod([subclass superclass], selector);
-        IMP originalImp = method_getImplementation(originalMethod);
-
-        IMP forwarderImp = [subclass instanceMethodForSelector:@selector(aMethodThatMustNotExist)];
-        class_addMethod(subclass, method_getName(originalMethod), forwarderImp, method_getTypeEncoding(originalMethod)); 
-
-        SEL aliasSelector = NSSelectorFromString([OCMRealMethodAliasPrefix stringByAppendingString:NSStringFromSelector(selector)]);
-        class_addMethod(subclass, aliasSelector, originalImp, method_getTypeEncoding(originalMethod));
-}
-
-- (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation
-{
-        // 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])];
-}
-
-
-
-#pragma mark  Overrides
-
-- (id)getNewRecorder
-{
-        return [[[OCPartialMockRecorder alloc] initWithSignatureResolver:self] autorelease];
+    [super addStub:aStub];
+    if(![aStub recordedAsClassMethod])
+        [self setupForwarderForSelector:[[aStub recordedInvocation] selector]];
 }
 
 - (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
 {
         [anInvocation invokeWithTarget:realObject];
 }
 
 
+#pragma mark  Subclass management
+
+- (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)setupForwarderForSelector:(SEL)sel
+{
+    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];
+}
+
+
+- (id)forwardingTargetForSelectorForRealObject:(SEL)sel
+{
+        // 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];
+}
+
+//  Make the compiler happy in -forwardingTargetForSelectorForRealObject: because it can't find the message…
+- (id)ocmock_replaced_forwardingTargetForSelector:(SEL)sel
+{
+    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];
+    }
+}
+
+
 @end