Use direct method interception via the runtime instead of OCMock

Author: ksperling-apple

src/darwin/Framework/CHIPTests/TestHelpers/MTRMockCB.m

@@ -17,8 +17,8 @@
 #import "MTRMockCB.h"
 
 #import <CoreBluetooth/CoreBluetooth.h>
-#import <OCMock/OCMock.h>
 #import <XCTest/XCTest.h>
+#import <objc/runtime.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -116,13 +116,49 @@ - (void)openL2CAPChannel:(CBL2CAPPSM)PSM;
 
 @implementation MTRMockCB {
 @package
-    id _classMock;
+    os_block_t _invalidate;
     dispatch_queue_t _queue;
     NSHashTable<MTRMockCBCentralManager *> * _managers;
     NSMutableDictionary<NSUUID *, MTRMockCBPeripheralDetails *> * _peripherals;
     CBManagerState _state;
 }
 
+static void InterceptClassMethod(__strong os_block_t * inOutCleanup, Class cls, SEL sel, id block)
+{
+    Method method = class_getClassMethod(cls, sel); // may return an inherited method
+    if (!method) {
+        NSString * reason = [NSString stringWithFormat:@"+[%@ %@] does not exist",
+                                      NSStringFromClass(cls), NSStringFromSelector(sel)];
+        @throw [NSException exceptionWithName:NSInternalInconsistencyException reason:reason userInfo:nil];
+    }
+    IMP originalImp = method_getImplementation(method);
+    if (imp_getBlock(originalImp)) {
+        NSString * reason = [NSString stringWithFormat:@"+[%@ %@] was already intercepted",
+                                      NSStringFromClass(cls), NSStringFromSelector(sel)];
+        @throw [NSException exceptionWithName:NSInternalInconsistencyException reason:reason userInfo:nil];
+    }
+
+    // Try to add the method first, in case it came from a super class.
+    // Note we need to pass the meta-class to class_addMethod().
+    IMP newImp = imp_implementationWithBlock(block);
+    if (class_addMethod(object_getClass(cls), sel, newImp, method_getTypeEncoding(method))) {
+        method = class_getClassMethod(cls, sel); // look up again so we clean up the method we added
+    } else {
+        method_setImplementation(method, newImp);
+    }
+
+    os_block_t nextCleanup = *inOutCleanup;
+    *inOutCleanup = ^{
+        // This isn't 100% correct if we added an override of a super-class method, because
+        // there is no API for removing a method. Instead we directly point it at the
+        // inherited implementation; this is good enough for our purposes here.
+        method_setImplementation(method, originalImp);
+        if (nextCleanup) {
+            nextCleanup();
+        }
+    };
+}
+
 - (instancetype)init
 {
     self = [super init];
@@ -131,29 +167,30 @@ - (instancetype)init
     _peripherals = [[NSMutableDictionary alloc] init];
     _state = CBManagerStatePoweredOn;
 
-    _classMock = OCMClassMock(CBCentralManager.class);
-    OCMStub(ClassMethod([_classMock alloc]))
-        .andCall(self, @selector(allocMockManager));
-    OCMStub(ClassMethod([_classMock authorization]))
-        .andCall(MTRMockCBCentralManager.class, @selector(authorization));
-    OCMStub(ClassMethod([_classMock supportsFeatures:0])).ignoringNonObjectArgs().andCall(MTRMockCBCentralManager.class, @selector(supportsFeatures:));
+    // Replace implementations of class methods we need to mock. We don't need to intercept
+    // any instance methods directly, because we're returning a mock object from `alloc`.
+    InterceptClassMethod(&_invalidate, CBCentralManager.class, @selector(alloc), ^id NS_RETURNS_RETAINED(void) {
+        return [[MTRMockCBCentralManager alloc] _initWithMock:self];
+    });
+    InterceptClassMethod(&_invalidate, CBCentralManager.class, @selector(supportsFeatures:), ^BOOL(CBCentralManagerFeature features) {
+        return [MTRMockCBCentralManager supportsFeatures:features];
+    });
+    InterceptClassMethod(&_invalidate, CBManager.class, @selector(authorization), ^CBManagerAuthorization(void) {
+        return [MTRMockCBCentralManager authorization];
+    });
+
     return self;
 }
 
 - (BOOL)isValid
 {
     __block BOOL valid;
     dispatch_sync(_queue, ^{
-        valid = (_classMock != nil);
+        valid = (_invalidate != nil);
     });
     return valid;
 }
 
-- (id)allocMockManager
-{
-    return [[MTRMockCBCentralManager alloc] _initWithMock:self];
-}
-
 - (void)reset
 {
     dispatch_sync(_queue, ^{
@@ -166,10 +203,11 @@ - (void)reset
 - (void)stopMocking
 {
     dispatch_sync(_queue, ^{
-        if (_classMock) {
+        if (_invalidate) {
+            _invalidate();
+            _invalidate = nil;
+
             NSArray<MTRMockCBCentralManager *> * managers = [_managers allObjects];
-            [_classMock stopMocking];
-            _classMock = nil;
             _managers = nil;
             _peripherals = nil;