Add the ability to ignore selectors on mocks to handle methods with variadic arguments.

Source/OCMock/OCClassMockObject.h

@@ -22,8 +22,8 @@
     Class originalMetaClass;
     Class classCreatedForNewMetaClass;
 }
-
-- (id)initWithClass:(Class)aClass;
+- (id)initForMockingClass:(Class)aClass;
+- (id)initForMockingInstancesOfClass:(Class)aClass;
 
 - (Class)mockedClass;
 - (Class)mockObjectClass; // since -class returns the mockedClass

Source/OCMock/OCClassMockObject.m

@@ -29,13 +29,20 @@ + (BOOL)supportsMocking:(NSString **)reason;
 @implementation OCClassMockObject
 
 #pragma mark Initialisers, description, accessors, etc.
-
-- (id)initWithClass:(Class)aClass
+- (id)initForMockingClass:(Class)aClass
+{
+    [self assertClassIsSupported:aClass];
+    [super init];
+    mockedClass = aClass;
+    [self prepareClassForMockingInstances:NO classes:YES];
+    return self;
+}
+- (id)initForMockingInstancesOfClass:(Class)aClass
 {
     [self assertClassIsSupported:aClass];
     [super init];
     mockedClass = aClass;
-    [self prepareClassForClassMethodMocking];
+    [self prepareClassForMockingInstances:YES classes:NO];
     return self;
 }
 
@@ -104,7 +111,7 @@ - (void)addStub:(OCMInvocationStub *)aStub
 
 #pragma mark Class method mocking
 
-- (void)prepareClassForClassMethodMocking
+- (void)prepareClassForMockingInstances:(BOOL)forInstances classes:(BOOL)forClasses
 {
     /* the runtime and OCMock depend on string and array; we don't intercept methods on them to avoid endless loops */
     if([[mockedClass class] isSubclassOfClass:[NSString class]] || [[mockedClass class] isSubclassOfClass:[NSArray class]])
@@ -116,56 +123,65 @@ - (void)prepareClassForClassMethodMocking
 
     /* if there is another mock for this exact class, stop it */
     id otherMock = OCMGetAssociatedMockForClass(mockedClass, NO);
-    if(otherMock != nil)
-        [otherMock stopMockingClassMethods];
-
-    OCMSetAssociatedMockForClass(self, mockedClass);
-
-    /* dynamically create a subclass and use its meta class as the meta class for the mocked class */
-    classCreatedForNewMetaClass = OCMCreateSubclass(mockedClass, mockedClass);
-    originalMetaClass = object_getClass(mockedClass);
-    id newMetaClass = object_getClass(classCreatedForNewMetaClass);
-
-    /* create a dummy initialize method */
-    Method myDummyInitializeMethod = class_getInstanceMethod([self mockObjectClass], @selector(initializeForClassObject));
-    const char *initializeTypes = method_getTypeEncoding(myDummyInitializeMethod);
-    IMP myDummyInitializeIMP = method_getImplementation(myDummyInitializeMethod);
-    class_addMethod(newMetaClass, @selector(initialize), myDummyInitializeIMP, initializeTypes);
-
-    object_setClass(mockedClass, newMetaClass); // only after dummy initialize is installed (iOS9)
-
-    /* point forwardInvocation: of the object to the implementation in the mock */
-    Method myForwardMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardInvocationForClassObject:));
-    IMP myForwardIMP = method_getImplementation(myForwardMethod);
-    class_addMethod(newMetaClass, @selector(forwardInvocation:), myForwardIMP, method_getTypeEncoding(myForwardMethod));
-
-    /* adding forwarder for most class methods (instance methods on meta class) to allow for verify after run */
-    NSArray *methodsNotToForward = @[
-        @"class", @"forwardingTargetForSelector:", @"methodSignatureForSelector:", @"forwardInvocation:", @"isBlock",
-        @"instanceMethodForwarderForSelector:", @"instanceMethodSignatureForSelector:", @"resolveClassMethod:"
-    ];
-    void (^setupForwarderFiltered)(Class, SEL) = ^(Class cls, SEL sel) {
-        if((cls == object_getClass([NSObject class])) || (cls == [NSObject class]) || (cls == object_getClass(cls)))
-            return;
-        if(OCMIsApplePrivateMethod(cls, sel))
-            return;
-        if([methodsNotToForward containsObject:NSStringFromSelector(sel)])
-            return;
-        @try
-        {
-            [self setupForwarderForClassMethodSelector:sel];
-        }
-        @catch(NSException *e)
-        {
-            // ignore for now
-        }
-    };
-    [NSObject enumerateMethodsInClass:originalMetaClass usingBlock:setupForwarderFiltered];
+    if( otherMock != nil && forClasses)
+    {
+        [NSException raise:NSInvalidArgumentException format:@"Class Mock must be made before partial mocks"];
+    }
+    if(otherMock == nil){
+        OCMSetAssociatedMockForClass(self, mockedClass);
+
+        /* dynamically create a subclass and use its meta class as the meta class for the mocked class */
+        classCreatedForNewMetaClass = OCMCreateSubclass(mockedClass, mockedClass);
+        originalMetaClass = object_getClass(mockedClass);
+        id newMetaClass = object_getClass(classCreatedForNewMetaClass);
+
+        /* create a dummy initialize method */
+        Method myDummyInitializeMethod = class_getInstanceMethod([self mockObjectClass], @selector(initializeForClassObject));
+        const char *initializeTypes = method_getTypeEncoding(myDummyInitializeMethod);
+        IMP myDummyInitializeIMP = method_getImplementation(myDummyInitializeMethod);
+        class_addMethod(newMetaClass, @selector(initialize), myDummyInitializeIMP, initializeTypes);
+
+        object_setClass(mockedClass, newMetaClass); // only after dummy initialize is installed (iOS9)
+
+        /* point forwardInvocation: of the object to the implementation in the mock */
+        Method myForwardMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardInvocationForClassObject:));
+        IMP myForwardIMP = method_getImplementation(myForwardMethod);
+        class_addMethod(newMetaClass, @selector(forwardInvocation:), myForwardIMP, method_getTypeEncoding(myForwardMethod));
+
+        /* adding forwarder for most class methods (instance methods on meta class) to allow for verify after run */
+        NSArray *methodsNotToForward = @[
+            @"class", @"forwardingTargetForSelector:", @"methodSignatureForSelector:", @"forwardInvocation:", @"isBlock",
+            @"instanceMethodForwarderForSelector:", @"instanceMethodSignatureForSelector:", @"resolveClassMethod:"
+        ];
+        void (^setupForwarderFiltered)(Class, SEL) = ^(Class cls, SEL sel) {
+            if((cls == object_getClass([NSObject class])) || (cls == [NSObject class]) || (cls == object_getClass(cls)))
+                return;
+            if(OCMIsApplePrivateMethod(cls, sel))
+                return;
+            if([methodsNotToForward containsObject:NSStringFromSelector(sel)])
+                return;
+            @try
+            {
+                [self setupForwarderForClassMethodSelector:sel];
+            }
+            @catch(NSException *e)
+            {
+                // ignore for now
+            }
+        };
+        [NSObject enumerateMethodsInClass:originalMetaClass usingBlock:setupForwarderFiltered];
+    }
 }
 
 
 - (void)setupForwarderForClassMethodSelector:(SEL)selector
 {
+    for(Class aClass in [OCMockObject classesIgnoringMockedSelector:selector])
+    {
+        if([mockedClass isKindOfClass:aClass])
+            return;
+    }
+
     SEL aliasSelector = OCMAliasForOriginalSelector(selector);
     if(class_getClassMethod(mockedClass, aliasSelector) != NULL)
         return;

Source/OCMock/OCMArg.h

@@ -36,6 +36,7 @@
 
 + (id *)setTo:(id)value;
 + (void *)setToValue:(NSValue *)value;
++ (id *)setToResultOfBlock:(id(^)(void))block;
 + (id)invokeBlock;
 + (id)invokeBlockWithArgs:(id)first, ... NS_REQUIRES_NIL_TERMINATION;
 

Source/OCMock/OCMArg.m

@@ -87,6 +87,10 @@ + (id *)setTo:(id)value
     return (id *)[[[OCMPassByRefSetter alloc] initWithValue:value] autorelease];
 }
 
++ (id *)setToResultOfBlock:(id(^)(void))block{
+    return (id *)[[[OCMPassByRefSetter alloc] initWithBlock:block] autorelease];
+}
+
 + (void *)setToValue:(NSValue *)value
 {
     return (id *)[[[OCMPassByRefSetter alloc] initWithValue:value] autorelease];

Source/OCMock/OCMPassByRefSetter.h

@@ -19,9 +19,11 @@
 @interface OCMPassByRefSetter : OCMArgAction
 {
     id value;
+    id(^block)(void);
 }
 
 - (id)initWithValue:(id)value;
+- (id)initWithBlock:(id(^)(void))block;
 
 + (BOOL)isPassByRefSetterInstance:(void *)ptr;
 

Source/OCMock/OCMPassByRefSetter.m

@@ -55,9 +55,23 @@ - (id)initWithValue:(id)aValue
     return self;
 }
 
+- (id)initWithBlock:(id(^)(void))aBlock{
+    if((self = [super init]))
+    {
+        block = [aBlock copy];
+        @synchronized(_OCMPassByRefSetterInstances)
+        {
+            NSHashInsertKnownAbsent(_OCMPassByRefSetterInstances, self);
+        }
+    }
+
+    return self;
+}
+
 - (void)dealloc
 {
     [value release];
+    [block release];
     @synchronized(_OCMPassByRefSetterInstances)
     {
         NSHashRemove(_OCMPassByRefSetterInstances, self);
@@ -70,10 +84,15 @@ - (void)handleArgument:(id)arg
     void *pointerValue = [arg pointerValue];
     if(pointerValue != NULL)
     {
-        if([value isKindOfClass:[NSValue class]])
-            [(NSValue *)value getValue:pointerValue];
-        else
-            *(id *)pointerValue = value;
+        if (block){
+           *(id*)pointerValue = block();
+        }
+        else{
+            if([value isKindOfClass:[NSValue class]])
+                [(NSValue *)value getValue:pointerValue];
+            else
+                *(id *)pointerValue = value;
+        }
     }
 }
 

Source/OCMock/OCMockMacros.h

@@ -155,3 +155,5 @@
     macro \
     _Pragma("clang diagnostic pop") \
 })
+
+#define OCMIgnore(class,selector) do { [OCMockObject  ignoreMethod:selector forClass:class];} while(0);

Source/OCMock/OCMockObject.h

@@ -43,6 +43,9 @@
 
 + (id)observerMock __deprecated_msg("Please use XCTNSNotificationExpectation instead.");
 
++ (NSArray <Class> *)classesIgnoringMockedSelector:(SEL)aSelector;
++ (void)ignoreMethod:(SEL) selector forClass:(Class) aClass;
+
 - (instancetype)init;
 
 - (void)setExpectationOrderMatters:(BOOL)flag;
@@ -62,6 +65,7 @@
 // internal use only
 
 - (void)addStub:(OCMInvocationStub *)aStub;
+- (void)removeStubsForSelector:(SEL)selector;
 - (void)addExpectation:(OCMInvocationExpectation *)anExpectation;
 - (void)addInvocation:(NSInvocation *)anInvocation;
 

Source/OCMock/OCMockObject.m

@@ -45,7 +45,7 @@ + (void)initialize
 
 + (id)mockForClass:(Class)aClass
 {
-    return [[[OCClassMockObject alloc] initWithClass:aClass] autorelease];
+    return [[[OCClassMockObject alloc] initForMockingClass:aClass] autorelease];
 }
 
 + (id)mockForProtocol:(Protocol *)aProtocol
@@ -83,6 +83,26 @@ + (id)observerMock
 }
 
 
+NSDictionary <NSString*,NSSet <Class> *>  * _static_ignoredBySelectorName;
+
++ (void)ignoreMethod:(SEL) aSelector forClass:(Class) aClass{
+    @synchronized (self) 
+    {
+        NSMutableDictionary * muIgnoreListBySelector = [[_static_ignoredBySelectorName?:@{} mutableCopy] autorelease];
+        NSMutableSet * muIgnoreList = [[muIgnoreListBySelector[NSStringFromSelector(aSelector)]?:NSSet.new mutableCopy] autorelease];
+        [muIgnoreList addObject:aClass ];
+        muIgnoreListBySelector[NSStringFromSelector(aSelector)] = [muIgnoreList.copy autorelease];
+        [_static_ignoredBySelectorName autorelease];
+        _static_ignoredBySelectorName = muIgnoreListBySelector.copy;
+    }
+}
+
++ (NSArray <Class> *)classesIgnoringMockedSelector:(SEL)aSelector{
+    return [_static_ignoredBySelectorName[NSStringFromSelector(aSelector)] allObjects];
+}
+
+
+
 #pragma mark Initialisers, description, accessors, etc.
 
 - (instancetype)init
@@ -190,6 +210,19 @@ - (void)setExpectationOrderMatters:(BOOL)flag
     expectationOrderMatters = flag;
 }
 
+- (void)removeStubsForSelector:(SEL)selector{
+    @synchronized (stubs) {
+        int s = stubs.count;
+        while (s--){
+            OCMInvocationStub * stub = stubs[s];
+            if ([stub isKindOfClass:OCMInvocationStub.class]){
+                if (stub.recordedInvocation.selector == selector){
+                    [stubs removeObjectAtIndex:s];
+                }
+            }
+        }
+    }
+}
 - (void)stopMocking
 {
     // invocations can contain objects that clients expect to be deallocated by now,
@@ -363,6 +396,20 @@ - (id)forwardingTargetForSelector:(SEL)aSelector
         [recorder setShouldReturnMockFromInit:(class_getInstanceMethod(object_getClass(recorder), aSelector) == NO)];
         return recorder;
     }
+    if (object_getClass(self) == OCPartialMockObject.class) 
+    {
+        Ivar ivar = class_getInstanceVariable(OCPartialMockObject.class,"realObject");
+        id mockObject = ivar?object_getIvar(self, ivar):nil;
+        if(mockObject)
+        {
+            for(Class aClass in [OCMockObject classesIgnoringMockedSelector:aSelector])
+            {
+                if([mockObject isKindOfClass:aClass])
+                    return mockObject;
+            }
+        }
+    }
+
     return nil;
 }
 

Source/OCMock/OCPartialMockObject.m

@@ -34,7 +34,7 @@ - (id)initWithObject:(NSObject *)anObject
     if([anObject isProxy])
         [NSException raise:NSInvalidArgumentException format:@"OCMock does not support partially mocking subclasses of NSProxy."];
     Class const class = [self classToSubclassForObject:anObject];
-    [super initWithClass:class];
+    [super initForMockingInstancesOfClass:class];
     realObject = [anObject retain];
     [self prepareObjectForInstanceMethodMocking];
     return self;
@@ -199,6 +199,12 @@ - (void)prepareObjectForInstanceMethodMocking
 
 - (void)setupForwarderForSelector:(SEL)sel
 {
+
+    for(Class aClass in [OCMockObject classesIgnoringMockedSelector:sel])
+    {
+        if([realObject isKindOfClass:aClass])
+            return;
+    }
     SEL aliasSelector = OCMAliasForOriginalSelector(sel);
     if(class_getInstanceMethod(object_getClass(realObject), aliasSelector) != NULL)
         return;
@@ -232,6 +238,14 @@ - (Class)classForRealObject
 - (id)forwardingTargetForSelectorForRealObject:(SEL)sel
 {
     // in here "self" is a reference to the real object, not the mock
+
+    for(Class aClass in [OCMockObject classesIgnoringMockedSelector:sel])
+    {
+        if([self isKindOfClass:aClass])
+            return self;
+    }
+
+
     OCPartialMockObject *mock = OCMGetAssociatedMockForObject(self);
     if(mock == nil)
         [NSException raise:NSInternalInconsistencyException format:@"No partial mock for object %p", self];