Make OCMInvocationConstraint retain invocation

Author: dmaclach

Source/OCMock/OCMConstraint.h

@@ -59,10 +59,12 @@
 
 @interface OCMInvocationConstraint : OCMConstraint
 {
-	@public
-	NSInvocation *invocation;
+  NSInvocation *invocation;
 }
 
+- (instancetype)initWithInvocation:(NSInvocation *)invocation NS_DESIGNATED_INITIALIZER;
+- (instancetype)init NS_UNAVAILABLE;
+
 @end
 
 @interface OCMBlockConstraint : OCMConstraint

Source/OCMock/OCMConstraint.m

@@ -15,7 +15,7 @@
  */
 
 #import <OCMock/OCMConstraint.h>
-
+#import "OCMFunctions.h"
 
 @implementation OCMConstraint
 
@@ -34,26 +34,30 @@ - (id)copyWithZone:(struct _NSZone *)zone
     return [self retain];
 }
 
++ (NSInvocation *)invocationWithSelector:(SEL)aSelector onObject:(id)anObject
+{
+  NSMethodSignature *signature = [anObject methodSignatureForSelector:aSelector];
+  if(signature == nil)
+    [NSException raise:NSInvalidArgumentException format:@"Unknown selector %@ used in constraint.", NSStringFromSelector(aSelector)];
+  NSInvocation *invocation = [NSInvocation invocationWithMethodSignature:signature];
+  [invocation setTarget:anObject];
+  [invocation setSelector:aSelector];
+  return invocation;
+}
+
 + (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject
 {
-	OCMInvocationConstraint *constraint = [OCMInvocationConstraint constraint];
-	NSMethodSignature *signature = [anObject methodSignatureForSelector:aSelector]; 
-	if(signature == nil)
-		[NSException raise:NSInvalidArgumentException format:@"Unkown selector %@ used in constraint.", NSStringFromSelector(aSelector)];
-	NSInvocation *invocation = [NSInvocation invocationWithMethodSignature:signature];
-	[invocation setTarget:anObject];
-	[invocation setSelector:aSelector];
-	constraint->invocation = invocation;
-	return constraint;
+  NSInvocation *invocation = [self invocationWithSelector:aSelector onObject:anObject];
+  return [[[OCMInvocationConstraint alloc] initWithInvocation:invocation] autorelease];
 }
 
 + (instancetype)constraintWithSelector:(SEL)aSelector onObject:(id)anObject withValue:(id)aValue
 {
-	OCMInvocationConstraint *constraint = (OCMInvocationConstraint *)[self constraintWithSelector:aSelector onObject:anObject];
-	if([[constraint->invocation methodSignature] numberOfArguments] < 4)
-		[NSException raise:NSInvalidArgumentException format:@"Constraint with value requires selector with two arguments."];
-	[constraint->invocation setArgument:&aValue atIndex:3];
-	return constraint;
+  NSInvocation *invocation = [self invocationWithSelector:aSelector onObject:anObject];
+  if([[invocation methodSignature] numberOfArguments] < 4)
+    [NSException raise:NSInvalidArgumentException format:@"Constraint with value requires selector with two arguments."];
+  [invocation setArgument:&aValue atIndex:3];
+  return [[[OCMInvocationConstraint alloc] initWithInvocation:invocation] autorelease];
 }
 
 
@@ -153,6 +157,33 @@ - (BOOL)evaluate:(id)value
 
 @implementation OCMInvocationConstraint
 
+- (instancetype)initWithInvocation:(NSInvocation *)anInvocation {
+  if((self = [super init]))
+  {
+    NSMethodSignature *signature = [anInvocation methodSignature];
+    if([signature numberOfArguments] < 3)
+    {
+      [NSException raise:NSInvalidArgumentException format:@"invocation must take at least one argument (other than _cmd and self)"];
+    }
+    if(!(OCMIsObjectType([signature getArgumentTypeAtIndex:2])))
+    {
+      [NSException raise:NSInvalidArgumentException format:@"invocation's second argument must be an object type"];      
+    }
+    if(strcmp([signature methodReturnType], @encode(BOOL)))
+    {
+      [NSException raise:NSInvalidArgumentException format:@"invocation must return BOOL"];      
+    }
+    invocation = [anInvocation retain];
+  }
+  return self;
+}
+
+- (void)dealloc
+{
+  [invocation release];
+  [super dealloc];
+}
+
 - (BOOL)evaluate:(id)value
 {
 	[invocation setArgument:&value atIndex:2]; // should test if constraint takes arg

Source/OCMockTests/OCMConstraintTests.m

@@ -205,4 +205,58 @@ - (void)testEvaluateNilBlockReturnsNo
     XCTAssertFalse([constraint evaluate:@"foo"]);
 }
 
+- (void)testEvaluateInvocationRetainsInvocation
+{
+  OCMInvocationConstraint *constraint;
+  @autoreleasepool {
+    SEL selector = @selector(checkArg:);
+    NSInvocation *anInvocation = [NSInvocation invocationWithMethodSignature:[self methodSignatureForSelector:selector]];
+    [anInvocation setTarget:self];
+    [anInvocation setSelector:selector];
+    constraint = [[OCMInvocationConstraint alloc] initWithInvocation:anInvocation];
+  }
+  XCTAssertTrue([constraint evaluate:@"foo"]);
+}
+
+- (BOOL)methodWithNoArgs
+{
+  return YES;
+}
+
+- (void)testEvaluateInvocationThrowsForInvocationForMethodWithoutArgument
+{
+  SEL selector = @selector(methodWithNoArgs);
+  NSInvocation *anInvocation = [NSInvocation invocationWithMethodSignature:[self methodSignatureForSelector:selector]];
+  [anInvocation setTarget:self];
+  [anInvocation setSelector:selector];
+  XCTAssertThrowsSpecificNamed([[OCMInvocationConstraint alloc] initWithInvocation:anInvocation], NSException, NSInvalidArgumentException);
+}
+
+- (BOOL)aMethodWithInt:(int)anInt
+{
+  return YES;
+}
+
+- (void)testEvaluateInvocationThrowsForInvocationForMethodWithoutObjectArgument
+{
+  SEL selector = @selector(aMethodWithInt:);
+  NSInvocation *anInvocation = [NSInvocation invocationWithMethodSignature:[self methodSignatureForSelector:selector]];
+  [anInvocation setTarget:self];
+  [anInvocation setSelector:selector];
+  XCTAssertThrowsSpecificNamed([[OCMInvocationConstraint alloc] initWithInvocation:anInvocation], NSException, NSInvalidArgumentException);
+}
+
+- (void)aMethodThatDoesNotReturnBool:(id)anArg
+{
+}
+
+- (void)testEvaluateInvocationThrowsForInvocationThatDoesNotReturnBool
+{
+  SEL selector = @selector(aMethodThatDoesNotReturnBool:);
+  NSInvocation *anInvocation = [NSInvocation invocationWithMethodSignature:[self methodSignatureForSelector:selector]];
+  [anInvocation setTarget:self];
+  [anInvocation setSelector:selector];
+  XCTAssertThrowsSpecificNamed([[OCMInvocationConstraint alloc] initWithInvocation:anInvocation], NSException, NSInvalidArgumentException);
+}
+
 @end