Merge pull request #128 from igorsales/master

Added support for stubbing singletons

README.md

@@ -160,6 +160,22 @@ __strong Class mockStringClass = mockClass([NSString class]);
 it explicitly strong as shown above, or use `id` instead.)
 
 
+How do you mock a singleton?
+-------------------------------
+
+```obj-c
+NSUserDefaults* defaults = mock([NSUserDefaults class]);
+__strong Class mockUserDefaultsClass = mockClass([NSUserDefaults class]);
+
+stubSingleton(mockUserDefaultsClass, standardUserDefaults);
+
+[given([NSUserDefaults standardUserDefaults]) willReturn:defaults];
+```
+
+(In the iOS 64-bit runtime, Class objects aren't strong by default. Either make
+it explicitly strong as shown above, or use `id` instead.)
+
+
 How do you mock a protocol?
 ---------------------------
 

Source/OCMockito/Core/OCMockito.h

@@ -192,6 +192,21 @@ FOUNDATION_EXPORT MKTOngoingStubbing *MKTGivenVoidWithLocation(id testCase, cons
 #endif
 
 
+#define MKTStubSingleton(mockClass, sel) \
+    [(MKTClassObjectMock*)mockClass swizzleSingletonAtSelector:@selector(sel)]
+
+#ifndef MKT_DISABLE_SHORT_SYNTAX
+/*!
+ * @abstract Stubs a singleton to the mock class object.
+ * @discussion
+ * <b>Name Clash</b><br />
+ * In the event of a name clash, <code>#define MKT_DISABLE_SHORT_SYNTAX</code> and use the synonym
+ * MKTStubSingleton instead.
+ */
+#define stubSingleton(mockClass, sel) MKTStubSingleton(mockClass, sel)
+#endif
+
+
 FOUNDATION_EXPORT id MKTVerifyWithLocation(id mock, id testCase, const char *fileName, int lineNumber);
 #define MKTVerify(mock) MKTVerifyWithLocation(mock, self, __FILE__, __LINE__)
 

Source/OCMockito/Mocking/MKTClassObjectMock.h

@@ -12,4 +12,7 @@
 
 - (instancetype)initWithClass:(Class)aClass;
 
+- (void)swizzleSingletonAtSelector:(SEL)singletonSelector;
+- (void)unswizzleSingletonAtSelector:(SEL)singletonSelector;
+
 @end

Source/OCMockito/Mocking/MKTClassObjectMock.m

@@ -1,16 +1,79 @@
 //  OCMockito by Jon Reid, http://qualitycoding.org/about/
 //  Copyright 2016 Jonathan M. Reid. See LICENSE.txt
 //  Contribution by David Hart
+//  Contribution by Igor Sales
 
 #import "MKTClassObjectMock.h"
+#import <objc/objc-runtime.h>
 
 
 @interface MKTClassObjectMock ()
 @property (nonatomic, strong, readonly) Class mockedClass;
 @end
 
+NSMutableDictionary* sSingletonMap = nil;
+
+@interface _MKTClassObjectMockMapEntry : NSObject {
+
+@public
+    __weak MKTClassObjectMock* _mock;
+
+}
+
+@property (nonatomic, weak,   readonly) MKTClassObjectMock* mock;
+@property (nonatomic, weak,   readonly) Class               mockedClass;
+@property (nonatomic, assign, readonly) IMP                 oldIMP;
+@property (nonatomic, assign, readonly) SEL                 selector;
+
+@end
+
+@implementation _MKTClassObjectMockMapEntry
+
+- (instancetype)initWithMock:(MKTClassObjectMock*)mock IMP:(IMP)oldIMP selector:(SEL)selector
+{
+    if (self = [super init]) {
+        _mock        = mock;
+        _mockedClass = mock.mockedClass;
+        _oldIMP      = oldIMP;
+        _selector    = selector;
+    }
+
+    return self;
+}
+
+@end
+
+
 @implementation MKTClassObjectMock
 
+
++ (void)initialize
+{
+    if (!sSingletonMap) {
+        sSingletonMap = [NSMutableDictionary new];
+    }
+}
+
+#define SINGLETON_KEY(C,S) [NSString stringWithFormat:@"%@-%@", C, NSStringFromSelector(S)]
+
++ (id)mockSingleton
+{
+    NSString* key = SINGLETON_KEY(self, _cmd);
+
+    _MKTClassObjectMockMapEntry* entry = sSingletonMap[key];
+
+    MKTClassObjectMock* mock = entry.mock;
+    if (mock) {
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
+        return [mock performSelector:_cmd withObject:nil];
+#pragma clang diagnostic pop
+    }
+
+    return nil;
+}
+
+
 - (instancetype)initWithClass:(Class)aClass
 {
     self = [super init];
@@ -19,6 +82,11 @@ - (instancetype)initWithClass:(Class)aClass
     return self;
 }
 
+- (void)dealloc {
+
+    [self unswizzleSingletons];
+}
+
 - (NSString *)description
 {
     return [@"mock class of " stringByAppendingString:NSStringFromClass(self.mockedClass)];
@@ -29,6 +97,77 @@ - (NSMethodSignature *)methodSignatureForSelector:(SEL)aSelector
     return [self.mockedClass methodSignatureForSelector:aSelector];
 }
 
+#pragma mark Operations
+
+- (void)swizzleSingletonAtSelector:(SEL)singletonSelector
+{
+    NSString* key = SINGLETON_KEY(_mockedClass, singletonSelector);
+
+    Method origMethod = class_getClassMethod(_mockedClass, singletonSelector);
+    Method newMethod  = class_getClassMethod([self class], @selector(mockSingleton));
+
+    IMP oldIMP = method_getImplementation(origMethod);
+    IMP newIMP = method_getImplementation(newMethod);
+
+    method_setImplementation(origMethod, newIMP);
+
+    _MKTClassObjectMockMapEntry* entry = sSingletonMap[key];
+    if (entry) {
+        // The user has already swizzled this singleton, keep the original implementation
+        oldIMP = entry.oldIMP;
+    }
+
+    sSingletonMap[key] = [[_MKTClassObjectMockMapEntry alloc] initWithMock:self
+                                                                       IMP:oldIMP
+                                                                  selector:singletonSelector];
+}
+
+- (void)unswizzleSingletonAtSelector:(SEL)singletonSelector {
+
+    NSString* key = SINGLETON_KEY(_mockedClass, singletonSelector);
+
+    _MKTClassObjectMockMapEntry* entry = sSingletonMap[key];
+
+    if (!entry) {
+        NSLog(@"Trying to unswizzle inexistent singleton method: + [%@ %@]",
+              self.mockedClass,
+              NSStringFromSelector(singletonSelector));
+        return;
+    }
+
+    [self unswizzleSingletonFromEntry:entry];
+}
+
+#pragma mark - Private
+
+- (void)unswizzleSingletonFromEntry:(_MKTClassObjectMockMapEntry*)swizzle
+{
+    NSAssert(swizzle, @"Invalid argument. swizzle argument cannot be nil");
+
+    Method origMethod = class_getClassMethod(swizzle.mockedClass, swizzle.selector);
+
+    method_setImplementation(origMethod, swizzle.oldIMP);
+}
+
+- (void)unswizzleSingletons {
+
+    NSMutableArray* keysToRemove = [NSMutableArray new];
+
+    [sSingletonMap enumerateKeysAndObjectsUsingBlock:^(NSString* key,
+                                                       _MKTClassObjectMockMapEntry* swizzle,
+                                                       BOOL* stop) {
+
+        //if (swizzle.mockedClass == self.mockedClass) {
+        // At time of dealloc, it's possible the weak ref to swizzle.mock is nil,
+        // so we also check directly on the struct member
+        if (swizzle.mock == self || swizzle->_mock == self) {
+            [self unswizzleSingletonFromEntry:swizzle];
+            [keysToRemove addObject:key];
+        }
+    }];
+
+    [sSingletonMap removeObjectsForKeys:keysToRemove];
+}
 
 #pragma mark NSObject protocol
 

Source/Tests/StubClassTests.m

@@ -6,13 +6,30 @@
 #import <OCHamcrest/OCHamcrest.h>
 #import <XCTest/XCTest.h>
 
+@interface MKTClassObjectMock()
+
++ (id)mockSingleton;
+
+@end
+
 @interface ClassMethodsReturningObject : NSObject
 @end
 
 @implementation ClassMethodsReturningObject
 
 + (id)methodReturningObject { return self; }
 
++ (id)singletonMethod
+{
+    static ClassMethodsReturningObject* sSingleton = nil;
+
+    if (!sSingleton) {
+        sSingleton = [ClassMethodsReturningObject new];
+    }
+
+    return sSingleton;
+}
+
 @end
 
 
@@ -37,4 +54,63 @@ - (void)testStubbedMethod_ShouldReturnGivenObject
     assertThat([myMockClass methodReturningObject], is(@"STUBBED"));
 }
 
+- (void)testStubbedSingleton_ShouldReturnGivenObject
+{
+    stubSingleton(myMockClass, singletonMethod);
+
+    [given([myMockClass singletonMethod]) willReturn:@"STUBBED"];
+
+    assertThat([ClassMethodsReturningObject singletonMethod], is(@"STUBBED"));
+}
+
+- (void)testStubbedSingletonOnExistingClass_ShouldReturnGivenObject
+{
+    Class userDefaultsClass = mockClass([NSUserDefaults class]);
+
+    stubSingleton(userDefaultsClass, standardUserDefaults);
+
+    [given([userDefaultsClass standardUserDefaults]) willReturn:@"STUBBED"];
+
+    assertThat([NSUserDefaults standardUserDefaults], is(@"STUBBED"));
+}
+
+- (void)testStubbedSingleton_LastSingletonStubTakesPrecedence
+{
+    stubSingleton(myMockClass, singletonMethod);
+
+    [given([myMockClass singletonMethod]) willReturn:@"STUBBED"];
+
+    Class myNewMockClass = mockClass([ClassMethodsReturningObject class]);
+
+    stubSingleton(myNewMockClass, singletonMethod);
+
+    [given([myNewMockClass singletonMethod]) willReturn:@"STUBBED2"];
+
+    assertThat([ClassMethodsReturningObject singletonMethod], is(@"STUBBED2"));
+}
+
+- (void)testStubbedSingleton_ValidUnswizzle
+{
+    stubSingleton(myMockClass, singletonMethod);
+
+    [given([myMockClass singletonMethod]) willReturn:@"STUBBED"];
+
+    MKTClassObjectMock* mock = (MKTClassObjectMock*)myMockClass;
+    [mock unswizzleSingletonAtSelector:@selector(singletonMethod)];
+
+    assertThat([ClassMethodsReturningObject singletonMethod], isNot(@"STUBBED"));
+}
+
+- (void)testStubbedSingleton_InvalidUnswizzling
+{
+    MKTClassObjectMock* mock = (MKTClassObjectMock*)myMockClass;
+
+    [mock unswizzleSingletonAtSelector:@selector(standardUserDefaults)];
+}
+
+- (void)testNoStubbedSingleton_ReturnsNil
+{
+    assertThat([MKTClassObjectMock mockSingleton], equalTo(nil));
+}
+
 @end