From 24ca2bb194314eabec4ed2f8bc2ed69ccdf72c16 Mon Sep 17 00:00:00 2001
From: Dave MacLachlan <dmaclach@gmail.com>
Date: Wed, 30 Dec 2020 16:04:35 -0800
Subject: [PATCH] Partial mocks use class created by OCClassMock

---
 Source/OCMock/OCPartialMockObject.m           | 19 +++----
 .../OCMockObjectPartialMocksTests.m           | 52 -------------------
 2 files changed, 10 insertions(+), 61 deletions(-)

diff --git a/Source/OCMock/OCPartialMockObject.m b/Source/OCMock/OCPartialMockObject.m
index 599b06e0..533d365a 100644
--- a/Source/OCMock/OCPartialMockObject.m
+++ b/Source/OCMock/OCPartialMockObject.m
@@ -86,12 +86,10 @@ - (void)stopMocking
 {
     if(realObject != nil)
     {
-        Class partialMockClass = object_getClass(realObject);
         OCMSetAssociatedMockForObject(nil, realObject);
         object_setClass(realObject, [self mockedClass]);
         [realObject release];
         realObject = nil;
-        OCMDisposeSubclass(partialMockClass);
     }
     [super stopMocking];
 }
@@ -150,27 +148,30 @@ - (void)prepareObjectForInstanceMethodMocking
 {
     OCMSetAssociatedMockForObject(self, realObject);
 
-    /* dynamically create a subclass and set it as the class of the object */
-    Class subclass = OCMCreateSubclass(mockedClass, realObject);
-    object_setClass(realObject, subclass);
+    if(!classCreatedForNewMetaClass) 
+    {
+        classCreatedForNewMetaClass = OCMCreateSubclass(mockedClass, mockedClass);
+    }
+    object_setClass(realObject, classCreatedForNewMetaClass);
 
     /* point forwardInvocation: of the object to the implementation in the mock */
 	Method myForwardMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardInvocationForRealObject:));
 	IMP myForwardIMP = method_getImplementation(myForwardMethod);
-    class_addMethod(subclass, @selector(forwardInvocation:), myForwardIMP, method_getTypeEncoding(myForwardMethod));
+    class_addMethod(classCreatedForNewMetaClass, @selector(forwardInvocation:), myForwardIMP, method_getTypeEncoding(myForwardMethod));
 
     /* do the same for forwardingTargetForSelector, remember existing imp with alias selector */
     Method myForwardingTargetMethod = class_getInstanceMethod([self mockObjectClass], @selector(forwardingTargetForSelectorForRealObject:));
     IMP myForwardingTargetIMP = method_getImplementation(myForwardingTargetMethod);
     IMP originalForwardingTargetIMP = [mockedClass instanceMethodForSelector:@selector(forwardingTargetForSelector:)];
-    class_addMethod(subclass, @selector(forwardingTargetForSelector:), myForwardingTargetIMP, method_getTypeEncoding(myForwardingTargetMethod));
-    class_addMethod(subclass, @selector(ocmock_replaced_forwardingTargetForSelector:), originalForwardingTargetIMP, method_getTypeEncoding(myForwardingTargetMethod));
+    const char *encoding = method_getTypeEncoding(myForwardingTargetMethod);
+    class_addMethod(classCreatedForNewMetaClass, @selector(forwardingTargetForSelector:), myForwardingTargetIMP, encoding);
+    class_addMethod(classCreatedForNewMetaClass, @selector(ocmock_replaced_forwardingTargetForSelector:), originalForwardingTargetIMP, encoding);
 
     /* We also override the -class method to return the original class */
     Method myObjectClassMethod = class_getInstanceMethod([self mockObjectClass], @selector(classForRealObject));
     const char *objectClassTypes = method_getTypeEncoding(myObjectClassMethod);
     IMP myObjectClassImp = method_getImplementation(myObjectClassMethod);
-    class_addMethod(subclass, @selector(class), myObjectClassImp, objectClassTypes);
+    class_addMethod(classCreatedForNewMetaClass, @selector(class), myObjectClassImp, objectClassTypes);
 
     /* Adding forwarder for most instance methods to allow for verify after run */
     NSArray *methodBlackList = @[@"class", @"forwardingTargetForSelector:", @"methodSignatureForSelector:", @"forwardInvocation:",
diff --git a/Source/OCMockTests/OCMockObjectPartialMocksTests.m b/Source/OCMockTests/OCMockObjectPartialMocksTests.m
index e9f3509e..a3487435 100644
--- a/Source/OCMockTests/OCMockObjectPartialMocksTests.m
+++ b/Source/OCMockTests/OCMockObjectPartialMocksTests.m
@@ -62,44 +62,6 @@ - (void)bar:(id)someArgument // maybe we should make it explicit that the arg is
 
 @end
 
-@interface TestClassThatObservesFoo : NSObject
-{
-    @public
-    id observedObject;
-}
-@end
-
-@implementation TestClassThatObservesFoo
-
-- (instancetype)initWithObject:(id)object
-{
-    if((self = [super init]))
-        observedObject = object;
-    return self;
-}
-
-- (void)dealloc
-{
-    [self stopObserving];
-}
-
-- (void)startObserving
-{
-    [observedObject addObserver:self forKeyPath:@"foo" options:0 context:NULL];
-}
-
-- (void)stopObserving
-{
-    if(observedObject != nil)
-    {
-        [observedObject removeObserver:self forKeyPath:@"foo" context:NULL];
-        observedObject = nil;
-    }
-}
-
-@end
-
-
 @interface TestClassThatCallsSelf : NSObject
 {
     int methodInt;
@@ -770,18 +732,4 @@ - (void)testDoesNotIncludeHintWhenStubbingIsNotGoingToHelp
     }
 }
 
-- (void)testThrowsExceptionWhenAttemptingToTearDownWrongClass
-{
-    TestClassWithSimpleMethod *realObject = [[TestClassWithSimpleMethod alloc] init];
-    TestClassThatObservesFoo *observer = [[TestClassThatObservesFoo alloc] initWithObject:realObject];
-    id mock = [OCMockObject partialMockForObject:realObject];
-    [observer startObserving];
-    
-    // If we invoked stopObserving here, then stopMocking would work; but we want to test the error case.
-    XCTAssertThrowsSpecificNamed([mock stopMocking], NSException, NSInvalidArgumentException);
-    
-    // Must reset the object here to avoid any attempt to remove the observer, which would fail.
-    observer->observedObject = nil;
-}
-
 @end