Using blocks to quickly transpose method implementations

BlockMock.h

#import <Foundation/Foundation.h>

typedef void (^BlockMockMethod)(NSInvocation *);

@interface BlockMock : NSProxy {
    /* Instance to use as the mock's default target */
    id _target;

    /* Map of NSString selectors to BlockMockMethod blocks */
    NSMutableDictionary *_mocks;
}

- (id) initWithTarget: (id) target;
- (void) mock: (SEL) selector withBlock: (BlockMockMethod) block;
- (void) unmock: (SEL) selector;

@end

BlockMock.m

#import "BlockMock.h"

/**
 * A block based object mocker.
 */
@implementation BlockMock

/**
 * Initialize with the provided default target. If a mock method is not available,
 * the method will be invoked on the target.
 */
- (id) initWithTarget: (id) target  {
    _target = [target retain];
    _mocks = [NSMutableDictionary new];

    return self;
}

- (void) dealloc {
    [_target release];
    [_mocks release];

    [super dealloc];
}

// NSProxy abstract method
- (void) forwardInvocation:(NSInvocation *)anInvocation {
    /* If available, use the mock implementation */
    BlockMockMethod m = [_mocks objectForKey: NSStringFromSelector([anInvocation selector])];

    if (m == NULL) {
        [anInvocation invokeWithTarget: _target];
        return;
    }

    m(anInvocation);
}

// NSProxy abstract method
- (NSMethodSignature *)methodSignatureForSelector:(SEL)aSelector {
    return [_target methodSignatureForSelector: aSelector];
}

/**
 * Attach a mock to the given method selector, setting the provided flag to YES if the
 * method is executed.
 */
- (void) flagOnRun: (SEL) selector flag: (BOOL *) flag {
    [self mock: selector withBlock: ^(NSInvocation *i) {
        *flag = YES;
        [i invoke];
    }];
}


/**
 * Add a mock block for the given selector.
 */
- (void) mock: (SEL) selector withBlock: (BlockMockMethod) block {
    [_mocks setObject: [[block copy] autorelease] forKey: NSStringFromSelector(selector)];
}

/**
 * Remove the registered mock for a given selector.
 */
- (void) unmock: (SEL) selector {
    [_mocks removeObjectForKey: NSStringFromSelector(selector)];
}

@end

BlockMockExample1.m

BlockMock *mock = [[BlockMock alloc] initWithTarget: @"Some String"];
NSString *fake = (NSString *) mock;

NSString *returnValue = @"I'm not the string you expected";
[mock mock: @selector(capitalizedString) withBlock: ^(NSInvocation *i) {
    [i setReturnValue: (void *) &returnValue];
}];

STAssertEqualObjects(returnValue, [fake capitalizedString], @"Wrong string");

BlockMockExample2.m

/* Verify that the connection is started */
BOOL connected = NO;

[_socket flagOnRun: @selector(connect:) flag: &connected];
[_client start];
STAssertTrue(connected && listening, @"Connection was not requested");