Implement stubAndReturn with different return values.

Bensk1 · May 2, 2012 · eff7f47 · eff7f47
1 parent 864aba4
commit eff7f47
Show file tree

Hide file tree

Showing 6 changed files with 107 additions and 5 deletions.
diff --git a/Kiwi/KWMock.h b/Kiwi/KWMock.h
@@ -65,8 +65,10 @@
 
 - (id)stub;
 - (id)stubAndReturn:(id)aValue;
+- (id)stubAndReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 - (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue;
+- (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 - (void)clearStubs;
 

diff --git a/Kiwi/KWMock.m b/Kiwi/KWMock.m
@@ -19,6 +19,8 @@
 static NSString * const StubTag = @"StubTag";
 static NSString * const ExpectTag = @"ExpectTag";
 static NSString * const StubValueKey = @"StubValueKey";
+static NSString * const StubSecondValueKey = @"StubSecondValueKey";
+static NSString * const ChangeStubValueAfterTimesKey = @"ChangeStubValueAfterTimesKey";
 
 @interface KWMock()
 
@@ -245,13 +247,25 @@ - (id)stubAndReturn:(id)aValue {
     return [KWInvocationCapturer invocationCapturerWithDelegate:self userInfo:userInfo];
 }
 
+- (id)stubAndReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {
+    NSDictionary *userInfo = [NSDictionary dictionaryWithObjectsAndKeys:StubTag, ExpectOrStubTagKey, aValue, StubValueKey, times, ChangeStubValueAfterTimesKey, aSecondValue, StubSecondValueKey, nil];
+    return [KWInvocationCapturer invocationCapturerWithDelegate:self userInfo:userInfo];
+}
+
 - (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue {
     [self expectMessagePattern:aMessagePattern];
     [self removeStubWithMessagePattern:aMessagePattern];
     KWStub *stub = [KWStub stubWithMessagePattern:aMessagePattern value:aValue];
     [self.stubs addObject:stub];
 }
 
+- (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {   
+    [self expectMessagePattern:aMessagePattern];
+    [self removeStubWithMessagePattern:aMessagePattern];
+    KWStub *stub = [KWStub stubWithMessagePattern:aMessagePattern value:aValue times:times afterThatReturn:aSecondValue];
+    [self.stubs addObject:stub];
+}
+
 - (void)clearStubs {
     [self.stubs removeAllObjects];
 }
@@ -321,10 +335,15 @@ - (NSMethodSignature *)invocationCapturer:(KWInvocationCapturer *)anInvocationCa
 - (void)invocationCapturer:(KWInvocationCapturer *)anInvocationCapturer didCaptureInvocation:(NSInvocation *)anInvocation {
     KWMessagePattern *messagePattern = [KWMessagePattern messagePatternFromInvocation:anInvocation];
     NSString *tag = [anInvocationCapturer.userInfo objectForKey:ExpectOrStubTagKey];
-
     if ([tag isEqualToString:StubTag]) {
         id value = [anInvocationCapturer.userInfo objectForKey:StubValueKey];
-        [self stubMessagePattern:messagePattern andReturn:value];
+        if (![anInvocationCapturer.userInfo objectForKey:StubSecondValueKey]) {
+            [self stubMessagePattern:messagePattern andReturn:value];
+        } else {
+            id times = [anInvocationCapturer.userInfo objectForKey:ChangeStubValueAfterTimesKey];
+            id secondValue = [anInvocationCapturer.userInfo objectForKey:StubSecondValueKey];
+            [self stubMessagePattern:messagePattern andReturn:value times:times afterThatReturn:secondValue];
+        }
     } else {
         [self expectMessagePattern:messagePattern];
     }

diff --git a/Kiwi/KWStub.h b/Kiwi/KWStub.h
@@ -12,22 +12,30 @@
 @private
     KWMessagePattern *messagePattern;
     id value;
+    id returnValueTimes;
+    int returnedValueTimes;
+    id secondValue;
 }
 
 #pragma mark -
 #pragma mark Initializing
 
 - (id)initWithMessagePattern:(KWMessagePattern *)aMessagePattern;
 - (id)initWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue;
+- (id)initWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 + (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern;
 + (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue;
++ (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 #pragma mark -
 #pragma mark Properties
 
 @property (nonatomic, readonly) KWMessagePattern *messagePattern;
 @property (nonatomic, readonly) id value;
+@property (nonatomic, readonly) id returnValueTimes;
+@property (nonatomic, readonly) int returnedValueTimes;
+@property (nonatomic, readonly) id secondValue;
 
 #pragma mark -
 #pragma mark Processing Invocations

diff --git a/Kiwi/KWStub.m b/Kiwi/KWStub.m
@@ -28,6 +28,17 @@ - (id)initWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValu
     return self;
 }
 
+- (id)initWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {
+    if ((self = [super init])) {
+        messagePattern = [aMessagePattern retain];
+        value = [aValue retain];
+        returnValueTimes = [times retain];
+        secondValue = [aSecondValue retain];
+    }
+
+    return self;
+}
+
 + (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern {
     return [self stubWithMessagePattern:aMessagePattern value:nil];
 }
@@ -36,9 +47,15 @@ + (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValu
     return [[[self alloc] initWithMessagePattern:aMessagePattern value:aValue] autorelease];
 }
 
++ (id)stubWithMessagePattern:(KWMessagePattern *)aMessagePattern value:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {
+    return [[[self alloc] initWithMessagePattern:aMessagePattern value:aValue times:times afterThatReturn:aSecondValue] autorelease];
+}
+
 - (void)dealloc {
     [messagePattern release];
     [value release];
+    [returnValueTimes release];
+    [secondValue release];
     [super dealloc];
 }
 
@@ -47,6 +64,9 @@ - (void)dealloc {
 
 @synthesize messagePattern;
 @synthesize value;
+@synthesize secondValue;
+@synthesize returnValueTimes;
+@synthesize returnedValueTimes;
 
 #pragma mark -
 #pragma mark Processing Invocations
@@ -81,11 +101,24 @@ - (void)writeWrappedValueToInvocationReturnValue:(NSInvocation *)anInvocation {
     const char *returnType = [[anInvocation methodSignature] methodReturnType];
     NSData *data = nil;
 
+    NSData *choosedForData = [self.value dataValue];
+
+    if (returnValueTimes != nil) {
+        NSString *returnValueTimesString = returnValueTimes;
+        int returnValueTimesInt = [returnValueTimesString intValue];
+
+        if (returnedValueTimes >= returnValueTimesInt) {
+            choosedForData = [self.secondValue dataValue];
+        }
+        returnedValueTimes++;
+    }
+
+
     // When the return type is not the same as the type of the wrapped value,
     // attempt to convert the wrapped value to the desired type.
 
     if (KWObjCTypeEqualToObjCType([self.value objCType], returnType))
-        data = [self.value dataValue];
+        data = choosedForData;
     else
         data = [self valueDataWithObjCType:returnType];
 
@@ -94,7 +127,20 @@ - (void)writeWrappedValueToInvocationReturnValue:(NSInvocation *)anInvocation {
 
 - (void)writeObjectValueToInvocationReturnValue:(NSInvocation *)anInvocation {
     assert(self.value && "self.value must not be nil");
-    [anInvocation setReturnValue:&value];
+
+    void *choosedForData = &value;
+
+    if (returnValueTimes != nil) {
+        NSString *returnValueTimesString = returnValueTimes;
+        int returnValueTimesInt = [returnValueTimesString intValue];
+
+        if (returnedValueTimes >= returnValueTimesInt) {
+            choosedForData = &secondValue;
+        }
+        returnedValueTimes++;
+    }
+
+    [anInvocation setReturnValue:choosedForData];
 
 #ifndef __clang_analyzer__
     NSString *selectorString = NSStringFromSelector([anInvocation selector]);

diff --git a/Kiwi/NSObject+KiwiStubAdditions.h b/Kiwi/NSObject+KiwiStubAdditions.h
@@ -22,8 +22,10 @@
 
 - (id)stub;
 - (id)stubAndReturn:(id)aValue;
+- (id)stubAndReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 - (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue;
+- (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue;
 
 - (void)clearStubs;
 

diff --git a/Kiwi/NSObject+KiwiStubAdditions.m b/Kiwi/NSObject+KiwiStubAdditions.m
@@ -13,6 +13,8 @@
 #import "KWStub.h"
 
 static NSString * const StubValueKey = @"StubValueKey";
+static NSString * const StubSecondValueKey = @"StubSecondValueKey";
+static NSString * const ChangeStubValueAfterTimesKey = @"ChangeStubValueAfterTimesKey";
 
 @implementation NSObject(KiwiStubAdditions)
 
@@ -32,7 +34,13 @@ - (NSMethodSignature *)invocationCapturer:(KWInvocationCapturer *)anInvocationCa
 - (void)invocationCapturer:(KWInvocationCapturer *)anInvocationCapturer didCaptureInvocation:(NSInvocation *)anInvocation {
     KWMessagePattern *messagePattern = [KWMessagePattern messagePatternFromInvocation:anInvocation];
     id value = [anInvocationCapturer.userInfo objectForKey:StubValueKey];
-    [self stubMessagePattern:messagePattern andReturn:value];
+    if (![anInvocationCapturer.userInfo objectForKey:StubSecondValueKey]) {
+        [self stubMessagePattern:messagePattern andReturn:value];
+    } else {
+        id times = [anInvocationCapturer.userInfo objectForKey:ChangeStubValueAfterTimesKey];
+        id secondValue = [anInvocationCapturer.userInfo objectForKey:StubSecondValueKey];
+        [self stubMessagePattern:messagePattern andReturn:value times:times afterThatReturn:secondValue];
+    }
 }
 
 #pragma mark -
@@ -71,6 +79,11 @@ - (id)stubAndReturn:(id)aValue {
     return [KWInvocationCapturer invocationCapturerWithDelegate:self userInfo:userInfo];
 }
 
+- (id)stubAndReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {
+    NSDictionary *userInfo = [NSDictionary dictionaryWithObjectsAndKeys: aValue, StubValueKey, times, ChangeStubValueAfterTimesKey, aSecondValue, StubSecondValueKey, nil];
+    return [KWInvocationCapturer invocationCapturerWithDelegate:self userInfo:userInfo];
+}
+
 - (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue {
     if ([self methodSignatureForSelector:aMessagePattern.selector] == nil) {
         [NSException raise:@"KWStubException" format:@"cannot stub -%@ because no such method exists",
@@ -83,6 +96,18 @@ - (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aVa
     KWAssociateObjectStub(self, stub);
 }
 
+- (void)stubMessagePattern:(KWMessagePattern *)aMessagePattern andReturn:(id)aValue times:(id)times afterThatReturn:(id)aSecondValue {   
+    if ([self methodSignatureForSelector:aMessagePattern.selector] == nil) {
+        [NSException raise:@"KWStubException" format:@"cannot stub -%@ because no such method exists",
+         NSStringFromSelector(aMessagePattern.selector)];
+    }
+
+    Class interceptClass = KWSetupObjectInterceptSupport(self);
+    KWSetupMethodInterceptSupport(interceptClass, aMessagePattern.selector);
+    KWStub *stub = [KWStub stubWithMessagePattern:aMessagePattern value:aValue times:times afterThatReturn:aSecondValue];
+    KWAssociateObjectStub(self, stub);
+}
+
 - (void)clearStubs {
     KWClearObjectStubs(self);
 }