Merge branch 'dmaclach-initLeak'

Author: erikdoe

Source/OCMock/NSInvocation+OCMAdditions.h

@@ -46,4 +46,11 @@
 - (NSString *)cStringDescriptionAtIndex:(NSInteger)anInt;
 - (NSString *)selectorDescriptionAtIndex:(NSInteger)anInt;
 
+- (BOOL)methodIsInInitFamily;
+- (BOOL)methodIsInAllocFamily;
+- (BOOL)methodIsInCopyFamily;
+- (BOOL)methodIsInMutableCopyFamily;
+- (BOOL)methodIsInNewFamily;
+
+
 @end

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -528,4 +528,59 @@ - (NSString *)selectorDescriptionAtIndex:(NSInteger)anInt
 	return [NSString stringWithFormat:@"@selector(%@)", NSStringFromSelector(selectorValue)];
 }
 
+
+- (BOOL)isMethodFamily:(NSString *)family
+{
+	// Definitions here: https://clang.llvm.org/docs/AutomaticReferenceCounting.html#method-families
+
+	NSMethodSignature *signature = [self methodSignature];
+	if(OCMIsObjectType(signature.methodReturnType) == NO)
+	{
+		return NO;
+	}
+
+	NSString *selString = NSStringFromSelector([self selector]);
+	NSRange underscoreRange = [selString rangeOfString:@"^_*" options:NSRegularExpressionSearch];
+	selString = [selString substringFromIndex:NSMaxRange(underscoreRange)];
+
+	if([selString hasPrefix:family] == NO)
+	{
+		return NO;
+	}
+	NSUInteger familyLength = [family length];
+	if(([selString length] > familyLength) &&
+			([[NSCharacterSet lowercaseLetterCharacterSet] characterIsMember:[selString characterAtIndex:familyLength]]))
+	{
+		return NO;
+	}
+	return YES;
+}
+
+
+- (BOOL)methodIsInInitFamily
+{
+	return [self isMethodFamily:@"init"];
+}
+
+- (BOOL)methodIsInAllocFamily
+{
+	return [self isMethodFamily:@"alloc"];
+}
+
+- (BOOL)methodIsInCopyFamily
+{
+	return [self isMethodFamily:@"copy"];
+}
+
+- (BOOL)methodIsInMutableCopyFamily
+{
+	return [self isMethodFamily:@"mutableCopy"];
+}
+
+- (BOOL)methodIsInNewFamily
+{
+	return [self isMethodFamily:@"new"];
+}
+
+
 @end

Source/OCMock/OCMFunctionsPrivate.h

@@ -40,3 +40,4 @@ void OCMSetAssociatedMockForObject(OCClassMockObject *mock, id anObject);
 OCPartialMockObject *OCMGetAssociatedMockForObject(id anObject);
 
 void OCMReportFailure(OCMLocation *loc, NSString *description);
+

Source/OCMock/OCMIndirectReturnValueProvider.m

@@ -14,9 +14,7 @@
  *  under the License.
  */
 
-#import "NSMethodSignature+OCMAdditions.h"
 #import "OCMIndirectReturnValueProvider.h"
-#import "NSInvocation+OCMAdditions.h"
 
 
 @implementation OCMIndirectReturnValueProvider

Source/OCMock/OCMInvocationMatcher.m

@@ -20,9 +20,7 @@
 #import "OCMPassByRefSetter.h"
 #import "NSInvocation+OCMAdditions.h"
 #import "OCMInvocationMatcher.h"
-#import "OCClassMockObject.h"
 #import "OCMFunctionsPrivate.h"
-#import "OCMBlockArgCaller.h"
 
 
 @interface NSObject(HCMatcherDummy)

Source/OCMock/OCMInvocationStub.m

@@ -15,11 +15,12 @@
  */
 
 #import "OCMInvocationStub.h"
-#import "OCMFunctionsPrivate.h"
 #import "OCMArg.h"
 #import "OCMArgAction.h"
 #import "NSInvocation+OCMAdditions.h"
 
+#define UNSET_RETURN_VALUE_MARKER ((id)0x01234567)
+
 @implementation OCMInvocationStub
 
 - (id)init
@@ -48,6 +49,29 @@ - (NSArray *)invocationActions
 
 
 - (void)handleInvocation:(NSInvocation *)anInvocation
+{
+    [self invokeArgActionsForInvocation:anInvocation];
+
+    if([anInvocation methodIsInInitFamily])
+    {
+        id returnVal = UNSET_RETURN_VALUE_MARKER;
+        [anInvocation setReturnValue:&returnVal];
+
+        [self invokeActionsForInvocation:anInvocation];
+
+        [anInvocation getReturnValue:&returnVal];
+        if(returnVal == UNSET_RETURN_VALUE_MARKER)
+        {
+            [NSException raise:NSInvalidArgumentException format:@"%@ was stubbed but no return value set. A return value is required for an init method. If you intended to return nil, make this explicit with .andReturn(nil)", NSStringFromSelector([anInvocation selector])];
+        }
+    }
+    else
+    {
+        [self invokeActionsForInvocation:anInvocation];
+    }
+}
+
+- (void)invokeArgActionsForInvocation:(NSInvocation *)anInvocation
 {
     NSMethodSignature *signature = [recordedInvocation methodSignature];
     NSUInteger n = [signature numberOfArguments];
@@ -62,13 +86,15 @@ - (void)handleInvocation:(NSInvocation *)anInvocation
         if([recordedArg isKindOfClass:[NSValue class]])
             recordedArg = [OCMArg resolveSpecialValues:recordedArg];
 
-        if(![recordedArg isKindOfClass:[OCMArgAction class]])
-            continue;
-
-        [recordedArg handleArgument:passedArg];
+        if([recordedArg isKindOfClass:[OCMArgAction class]])
+            [recordedArg handleArgument:passedArg];
     }
+}
 
+- (void)invokeActionsForInvocation:(NSInvocation *)anInvocation
+{
     [invocationActions makeObjectsPerformSelector:@selector(handleInvocation:) withObject:anInvocation];
 }
 
+
 @end

Source/OCMock/OCMMacroState.m

@@ -16,10 +16,7 @@
 
 #import "OCMMacroState.h"
 #import "OCMStubRecorder.h"
-#import "OCMockObject.h"
 #import "OCMExpectationRecorder.h"
-#import "OCMVerifier.h"
-#import "OCMInvocationMatcher.h"
 
 
 @implementation OCMMacroState
@@ -71,14 +68,17 @@ + (OCMStubRecorder *)endExpectMacro
 + (void)beginRejectMacro
 {
     OCMExpectationRecorder *recorder = [[[OCMExpectationRecorder alloc] init] autorelease];
-    [recorder never];
     OCMMacroState *macroState = [[OCMMacroState alloc] initWithRecorder:recorder];
     [NSThread currentThread].threadDictionary[OCMGlobalStateKey] = macroState;
     [macroState release];
 }
 
 + (OCMStubRecorder *)endRejectMacro
 {
+    OCMMacroState *globalState = [NSThread currentThread].threadDictionary[OCMGlobalStateKey];
+    // Calling never after the invocation to avoid running afoul of ARC's expectations on
+    // return values from init methods.
+    [(OCMExpectationRecorder *)[globalState recorder] never];
     return [self endStubMacro];
 }
 

Source/OCMock/OCMNonRetainingObjectReturnValueProvider.m

@@ -16,7 +16,7 @@
 
 #import "OCMFunctions.h"
 #import "OCMNonRetainingObjectReturnValueProvider.h"
-
+#import "NSInvocation+OCMAdditions.h"
 
 @implementation OCMNonRetainingObjectReturnValueProvider
 
@@ -33,12 +33,19 @@ - (void)handleInvocation:(NSInvocation *)anInvocation
     {
         @throw [NSException exceptionWithName:NSInvalidArgumentException reason:@"Expected invocation with object return type. Did you mean to use andReturnValue: instead?" userInfo:nil];
     }
-    NSString *sel = NSStringFromSelector([anInvocation selector]);
-    if([sel hasPrefix:@"alloc"] || [sel hasPrefix:@"new"] || [sel hasPrefix:@"copy"] || [sel hasPrefix:@"mutableCopy"])
+
+    if([anInvocation methodIsInAllocFamily] || [anInvocation methodIsInNewFamily] ||
+            [anInvocation methodIsInCopyFamily] || [anInvocation methodIsInMutableCopyFamily])
     {
         // methods that "create" an object return it with an extra retain count
         [returnValue retain];
     }
+    else if([anInvocation methodIsInInitFamily])
+    {
+        // init family methods "consume" self and retain their return value. Do the retain first in case the return value and self are the same.
+        [returnValue retain];
+        [[anInvocation target] release];
+    }
     [anInvocation setReturnValue:&returnValue];
 }
 @end

Source/OCMock/OCMQuantifier.m

@@ -14,10 +14,9 @@
  *  under the License.
  */
 
-#import <Foundation/Foundation.h>
+#import "OCMQuantifier.h"
 #import "OCMMacroState.h"
 #import "OCMVerifier.h"
-#import "OCMQuantifier.h"
 
 
 @interface OCMExactCountQuantifier : OCMQuantifier

Source/OCMock/OCMRecorder.h

@@ -25,19 +25,23 @@
     OCMockObject         *mockObject;
     OCMInvocationMatcher *invocationMatcher;
     BOOL                 wasUsed;
+    BOOL                 shouldReturnMockFromInit;
 }
 
+
 - (instancetype)init;
 - (instancetype)initWithMockObject:(OCMockObject *)aMockObject;
 
 - (void)setMockObject:(OCMockObject *)aMockObject;
+- (void)setShouldReturnMockFromInit:(BOOL)flag;
 
 - (OCMInvocationMatcher *)invocationMatcher;
 - (BOOL)wasUsed;
 
 - (id)classMethod;
 - (id)ignoringNonObjectArgs;
 
+
 @end
 
 @interface OCMRecorder (Properties)