From 17e9ee893393c24153fbbc8bc5a69fdcd22a393c Mon Sep 17 00:00:00 2001
From: Jack Wu <jkwu@fb.com>
Date: Mon, 29 Jan 2018 15:42:07 -0500
Subject: [PATCH] Retain class type arguments of NSInvocation

---
 Source/OCMock/NSInvocation+OCMAdditions.m |  4 ++--
 Source/OCMock/OCMFunctions.m              |  5 -----
 Source/OCMock/OCMFunctionsPrivate.h       |  1 -
 Source/OCMockTests/OCMockObjectTests.m    | 26 ++++++++++++++++++++++-
 4 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/Source/OCMock/NSInvocation+OCMAdditions.m b/Source/OCMock/NSInvocation+OCMAdditions.m
index 112d7a77..ae5f1c38 100644
--- a/Source/OCMock/NSInvocation+OCMAdditions.m
+++ b/Source/OCMock/NSInvocation+OCMAdditions.m
@@ -80,7 +80,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     for(NSUInteger index = 2; index < numberOfArguments; index++)
     {
         const char *argumentType = [[self methodSignature] getArgumentTypeAtIndex:index];
-        if(OCMIsObjectType(argumentType) && !OCMIsClassType(argumentType))
+        if(OCMIsObjectType(argumentType))
         {
             id argument;
             [self getArgument:&argument atIndex:index];
@@ -102,7 +102,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     }
 
     const char *returnType = [[self methodSignature] methodReturnType];
-    if(OCMIsObjectType(returnType) && !OCMIsClassType(returnType))
+    if(OCMIsObjectType(returnType))
     {
         id returnValue;
         [self getReturnValue:&returnValue];
diff --git a/Source/OCMock/OCMFunctions.m b/Source/OCMock/OCMFunctions.m
index 758ed3d2..69bbdaad 100644
--- a/Source/OCMock/OCMFunctions.m
+++ b/Source/OCMock/OCMFunctions.m
@@ -48,11 +48,6 @@ static BOOL OCMIsUnqualifiedClassType(const char *unqualifiedObjCType)
     return (strcmp(unqualifiedObjCType, @encode(Class)) == 0);
 }
 
-BOOL OCMIsClassType(const char *objCType)
-{
-    return OCMIsUnqualifiedClassType(OCMTypeWithoutQualifiers(objCType));
-}
-
 
 static BOOL OCMIsUnqualifiedBlockType(const char *unqualifiedObjCType)
 {
diff --git a/Source/OCMock/OCMFunctionsPrivate.h b/Source/OCMock/OCMFunctionsPrivate.h
index 1984c228..96a2721a 100644
--- a/Source/OCMock/OCMFunctionsPrivate.h
+++ b/Source/OCMock/OCMFunctionsPrivate.h
@@ -21,7 +21,6 @@
 @class OCPartialMockObject;
 
 
-BOOL OCMIsClassType(const char *objCType);
 BOOL OCMIsBlockType(const char *objCType);
 BOOL OCMIsObjectType(const char *objCType);
 const char *OCMTypeWithoutQualifiers(const char *objCType);
diff --git a/Source/OCMockTests/OCMockObjectTests.m b/Source/OCMockTests/OCMockObjectTests.m
index af015e86..b8457009 100644
--- a/Source/OCMockTests/OCMockObjectTests.m
+++ b/Source/OCMockTests/OCMockObjectTests.m
@@ -171,6 +171,21 @@ - (void)doStuffWithBlock:(void (^)(id<TestProtocol> arg))block;
 
 @end
 
+@interface TestClassWithClassArgMethod : NSObject
+
+- (void)doStuffWithClass:(Class)aClass;
+
+@end
+
+@implementation TestClassWithClassArgMethod
+
+- (void)doStuffWithClass:(Class)aClass
+{
+    // stubbed out anyway
+}
+
+@end
+
 static NSString *TestNotification = @"TestNotification";
 
 
@@ -985,7 +1000,6 @@ - (void)testFailsVerifyExpectedMethodsWithDelay
 	XCTAssertThrows([mock verifyWithDelay:0.1], @"Should have raised an exception because method was not called.");
 }
 
-
 // --------------------------------------------------------------------------------------
 //	ordered expectations
 // --------------------------------------------------------------------------------------
@@ -1170,6 +1184,16 @@ - (void)testVerifyWithDelayDoesNotWaitForRejects
     XCTAssertTrue([end timeIntervalSinceDate:start] < 3, @"Should have returned before delay was up");
 }
 
+- (void)testClassArgsAreRetained
+{
+
+    id mockWithClassMethod = [OCMockObject mockForClass:[TestClassWithClassArgMethod class]];
+    @autoreleasepool {
+        [[mockWithClassMethod stub] doStuffWithClass:[OCMArg any]];
+    }
+    XCTAssertNoThrow([mockWithClassMethod doStuffWithClass:[NSString class]]);
+}
+
 @end