Allow capturing of nil arguments

agilebits · Oct 10, 2012 · c1a1a0b · c1a1a0b
1 parent 4b6860b
commit c1a1a0b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/Kiwi/KWCaptureSpy.h b/Kiwi/KWCaptureSpy.h
@@ -3,6 +3,7 @@
 #import "KWMessageSpying.h"
 
 @interface KWCaptureSpy : NSObject<KWMessageSpying> {
+    BOOL _argumentCaptured;
     id _argument;
     NSUInteger _argumentIndex;
 }

diff --git a/Kiwi/KWCaptureSpy.m b/Kiwi/KWCaptureSpy.m
@@ -10,19 +10,20 @@ @implementation KWCaptureSpy
 - (id)initWithArgumentIndex:(NSUInteger)index {
     if ((self = [super init])) {
         _argumentIndex = index;
+        _argumentCaptured = NO;
     }
     return self;
 }
 
 - (id)argument {
-    if (!_argument) {
+    if (!_argumentCaptured) {
         @throw [NSException exceptionWithName:NSInternalInconsistencyException reason:@"Argument requested has yet to be captured." userInfo:nil];    
     }
     return [[_argument retain] autorelease];
 }
 
 - (void)object:(id)anObject didReceiveInvocation:(NSInvocation *)anInvocation {
-    if (!_argument) {
+    if (!_argumentCaptured) {
         NSMethodSignature *signature = [anInvocation methodSignature];
         const char *objCType = [signature messageArgumentTypeAtIndex:_argumentIndex];        
         if (KWObjCTypeIsObject(objCType)) {
@@ -37,6 +38,7 @@ - (void)object:(id)anObject didReceiveInvocation:(NSInvocation *)anInvocation {
             NSData *data = [anInvocation messageArgumentDataAtIndex:_argumentIndex];
             _argument = [[KWValue valueWithBytes:[data bytes] objCType:objCType] retain];
         }
+        _argumentCaptured = YES;
     }
 }
 

diff --git a/Tests/KWCaptureTest.m b/Tests/KWCaptureTest.m
@@ -52,6 +52,15 @@ - (void)testShouldBeAbleToCaptureValues {
     STAssertEqualObjects(spy.argument, [KWValue valueWithDouble:2], @"Captured argument should be equal to '2'");        
 }
 
+- (void)testShouldBeAbleToCaptureNils {
+    id robotMock = [KWMock nullMockForClass:[Robot class]];
+    KWCaptureSpy *spy = [robotMock captureArgument:@selector(speak:afterDelay:whenDone:) atIndex:0];
+
+    [robotMock speak:nil afterDelay:2 whenDone:^{}];
+
+    STAssertNil(spy.argument, @"Captured argument should be nil");
+}
+
 - (void)testShouldRaiseAnExceptionIfArgumentHasNotBeenCaptured {
     id robotMock = [KWMock nullMockForClass:[Robot class]];
     KWCaptureSpy *spy = [robotMock captureArgument:@selector(speak:afterDelay:whenDone:) atIndex:1];