Test update OCMock

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"
-
-
-@interface OCPartialMockObject (Private)
-- (void)forwardInvocationForRealObject:(NSInvocation *)anInvocation;
-@end 
-
-
-NSString *OCMRealMethodAliasPrefix = @"ocmock_replaced_";
+#import "NSMethodSignature+OCMAdditions.h"
+#import "NSObject+OCMAdditions.h"
+#import "OCMFunctions.h"
+#import "OCMInvocationStub.h"
 
 @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
 {
-        [super initWithClass:[anObject class]];
-        realObject = [anObject retain];
-        [[self class] rememberPartialMock:self forObject:anObject];
-        [self setupSubclassForObject:realObject];
-        return self;
+  NSParameterAssert(anObject != nil);
+  [self assertClassIsSupported:[anObject class]];
+  [super initWithClass:[anObject class]];
+  realObject = [anObject retain];
+  [self prepareObjectForInstanceMethodMocking];
+  return self;
 }
 
 - (void)dealloc
 {
-        if(realObject != nil)
-                [self stop];
-        [super dealloc];
+  [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
-{
-        object_setClass(realObject, [self mockedClass]);
-        [realObject release];
-        [[self class] forgetPartialMockForObject:realObject];
-        realObject = nil;
+#pragma mark Helper methods
+
+- (void)assertClassIsSupported:(Class) class {
+  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 Extending/overriding superclass behaviour
+
+    - (void)stopMocking {
+  if (realObject != nil) {
+    OCMSetAssociatedMockForObject(nil, realObject);
+    object_setClass(realObject, [self mockedClass]);
+    [realObject release];
+    realObject = nil;
+  }
+  [super stopMocking];
+}
+
+- (void)addStub:(OCMInvocationStub*)aStub {
+  [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];
+  }
 }
 
 
-#pragma mark  Subclass management
-
-- (void)setupSubclassForObject:(id)anObject
-{
-        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);
-}
-
-- (void)setupForwarderForSelector:(SEL)selector
-{
-        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];
-}
-
-- (void)handleUnRecordedInvocation:(NSInvocation *)anInvocation
-{
-        [anInvocation invokeWithTarget:realObject];
-}
-
-
 @end