Fix memory leaks, add section-object support to new test harness (#360)

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -370,8 +370,8 @@
 		CCA282CD1E9EB73E0037E8B7 /* ASTipNode.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA282CB1E9EB73E0037E8B7 /* ASTipNode.m */; };
 		CCA282D01E9EBF6C0037E8B7 /* ASTipsWindow.h in Headers */ = {isa = PBXBuildFile; fileRef = CCA282CE1E9EBF6C0037E8B7 /* ASTipsWindow.h */; };
 		CCA282D11E9EBF6C0037E8B7 /* ASTipsWindow.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA282CF1E9EBF6C0037E8B7 /* ASTipsWindow.m */; };
-		CCA5F62E1EECC2A80060C137 /* ASAssert.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62D1EECC2A80060C137 /* ASAssert.m */; };
 		CCA5F62C1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62B1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m */; };
+		CCA5F62E1EECC2A80060C137 /* ASAssert.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62D1EECC2A80060C137 /* ASAssert.m */; };
 		CCB2F34D1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB2F34C1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m */; };
 		CCB338E41EEE11160081F21A /* OCMockObject+ASAdditions.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB338E31EEE11160081F21A /* OCMockObject+ASAdditions.m */; };
 		CCB338E71EEE27760081F21A /* ASTestCase.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB338E61EEE27760081F21A /* ASTestCase.m */; };
@@ -832,9 +832,9 @@
 		CCA282CB1E9EB73E0037E8B7 /* ASTipNode.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASTipNode.m; sourceTree = "<group>"; };
 		CCA282CE1E9EBF6C0037E8B7 /* ASTipsWindow.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASTipsWindow.h; sourceTree = "<group>"; };
 		CCA282CF1E9EBF6C0037E8B7 /* ASTipsWindow.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASTipsWindow.m; sourceTree = "<group>"; };
-		CCA5F62D1EECC2A80060C137 /* ASAssert.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASAssert.m; sourceTree = "<group>"; };
 		CCA5F62A1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "NSInvocation+ASTestHelpers.h"; sourceTree = "<group>"; };
 		CCA5F62B1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "NSInvocation+ASTestHelpers.m"; sourceTree = "<group>"; };
+		CCA5F62D1EECC2A80060C137 /* ASAssert.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASAssert.m; sourceTree = "<group>"; };
 		CCB2F34C1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASDisplayNodeSnapshotTests.m; sourceTree = "<group>"; };
 		CCB338E21EEE11160081F21A /* OCMockObject+ASAdditions.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "OCMockObject+ASAdditions.h"; sourceTree = "<group>"; };
 		CCB338E31EEE11160081F21A /* OCMockObject+ASAdditions.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "OCMockObject+ASAdditions.m"; sourceTree = "<group>"; };

Source/ASRunLoopQueue.h

@@ -51,6 +51,8 @@ AS_SUBCLASSING_RESTRICTED
 
 + (instancetype)sharedDeallocationQueue;
 
+- (void)test_drain;
+
 - (void)releaseObjectInBackground:(id)object;
 
 @end

Source/ASRunLoopQueue.mm

@@ -143,6 +143,29 @@ static void runLoopSourceCallback(void *info) {
   _thread = nil;
 }
 
+- (void)test_drain
+{
+  [self performSelector:@selector(_test_drain) onThread:_thread withObject:nil waitUntilDone:YES];
+}
+
+- (void)_test_drain
+{
+  while (true) {
+    @autoreleasepool {
+      _queueLock.lock();
+      std::deque<id> currentQueue = _queue;
+      _queue = std::deque<id>();
+      _queueLock.unlock();
+
+      if (currentQueue.empty()) {
+        return;
+      } else {
+        currentQueue.clear();
+      }
+    }
+  }
+}
+
 - (void)_stop
 {
   CFRunLoopStop(CFRunLoopGetCurrent());

Source/Details/ASWeakSet.m

@@ -27,7 +27,7 @@
 {
   self = [super init];
   if (self) {
-    _hashTable = [NSHashTable hashTableWithOptions:NSPointerFunctionsWeakMemory | NSPointerFunctionsObjectPointerPersonality];
+    _hashTable = [NSHashTable hashTableWithOptions:NSHashTableWeakMemory | NSHashTableObjectPointerPersonality];
   }
   return self;
 }

Tests/ASCollectionModernDataSourceTests.m

@@ -23,21 +23,28 @@
 @interface ASTestCellNode : ASCellNode
 @end
 
+@interface ASTestSection : NSObject <ASSectionContext>
+@property (nonatomic, readonly) NSMutableArray *viewModels;
+@end
+
 @implementation ASCollectionModernDataSourceTests {
 @private
   id mockDataSource;
   UIWindow *window;
   UIViewController *viewController;
   ASCollectionNode *collectionNode;
-  NSMutableArray<NSMutableArray *> *sections;
+  NSMutableArray<ASTestSection *> *sections;
 }
 
 - (void)setUp {
   [super setUp];
   // Default is 2 sections: 2 items in first, 1 item in second.
   sections = [NSMutableArray array];
-  [sections addObject:[NSMutableArray arrayWithObjects:[NSObject new], [NSObject new], nil]];
+  [sections addObject:[ASTestSection new]];
-  [sections addObject:[NSMutableArray arrayWithObjects:[NSObject new], nil]];
+  [sections[0].viewModels addObject:[NSObject new]];
+  [sections[0].viewModels addObject:[NSObject new]];
+  [sections addObject:[ASTestSection new]];
+  [sections[1].viewModels addObject:[NSObject new]];
   window = [[UIWindow alloc] initWithFrame:CGRectMake(0, 0, 100, 100)];
   viewController = [[UIViewController alloc] init];
 
@@ -54,6 +61,7 @@
    @selector(collectionNode:numberOfItemsInSection:),
    @selector(collectionNode:nodeBlockForItemAtIndexPath:),
    @selector(collectionNode:viewModelForItemAtIndexPath:),
+   @selector(collectionNode:contextForSection:),
    nil];
   [mockDataSource setExpectationOrderMatters:YES];
   
@@ -64,7 +72,6 @@
 - (void)tearDown
 {
   [collectionNode waitUntilAllUpdatesAreCommitted];
-  OCMVerifyAll(mockDataSource);
   [super tearDown];
 }
 
@@ -82,11 +89,12 @@
   // Reload at (0, 0)
   NSIndexPath *reloadedPath = [NSIndexPath indexPathForItem:0 inSection:0];
 
-  [self performUpdateReloadingItems:@{ reloadedPath: [NSObject new] }
+  [self performUpdateReloadingSections:nil
-                     reloadMappings:@{ reloadedPath: reloadedPath }
+                        reloadingItems:@{ reloadedPath: [NSObject new] }
-                     insertingItems:nil
+                        reloadMappings:@{ reloadedPath: reloadedPath }
-                      deletingItems:nil
+                        insertingItems:nil
-            skippedReloadIndexPaths:nil];
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 - (void)testInsertingAnItem
@@ -96,11 +104,12 @@
   // Insert at (1, 0)
   NSIndexPath *insertedPath = [NSIndexPath indexPathForItem:0 inSection:1];
 
-  [self performUpdateReloadingItems:nil
+  [self performUpdateReloadingSections:nil
-                     reloadMappings:nil
+                        reloadingItems:nil
-                     insertingItems:@{ insertedPath: [NSObject new] }
+                        reloadMappings:nil
-                      deletingItems:nil
+                        insertingItems:@{ insertedPath: [NSObject new] }
-            skippedReloadIndexPaths:nil];
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 - (void)testReloadingAnItemWithACompatibleViewModel
@@ -115,17 +124,27 @@
 
   // Cell node should get -canUpdateToViewModel:
   id mockCellNode = [collectionNode nodeForItemAtIndexPath:reloadedPath];
-  [mockCellNode setExpectationOrderMatters:YES];
   OCMExpect([mockCellNode canUpdateToViewModel:viewModel])
   .andReturn(YES);
 
-  [self performUpdateReloadingItems:@{ reloadedPath: viewModel }
+  [self performUpdateReloadingSections:nil
-                     reloadMappings:@{ reloadedPath: [NSIndexPath indexPathForItem:0 inSection:0] }
+                        reloadingItems:@{ reloadedPath: viewModel }
-                     insertingItems:nil
+                        reloadMappings:@{ reloadedPath: [NSIndexPath indexPathForItem:0 inSection:0] }
-                      deletingItems:@[ deletedPath ]
+                        insertingItems:nil
-            skippedReloadIndexPaths:@[ reloadedPath ]];
+                         deletingItems:@[ deletedPath ]
+               skippedReloadIndexPaths:@[ reloadedPath ]];
+}
 
-  OCMVerifyAll(mockCellNode);
+- (void)testReloadingASection
+{
+  [self loadInitialData];
+
+  [self performUpdateReloadingSections:@{ @0: [ASTestSection new] }
+                        reloadingItems:nil
+                        reloadMappings:nil
+                        insertingItems:nil
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 #pragma mark - Helpers
@@ -137,14 +156,19 @@
     // It reads all the counts
     [self expectDataSourceCountMethods];
 
+    // It reads each section object.
+    for (NSInteger section = 0; section < sections.count; section++) {
+      [self expectContextMethodForSection:section];
+    }
+
     // It reads the contents for each item.
     for (NSInteger section = 0; section < sections.count; section++) {
-      NSArray *items = sections[section];
+      NSArray *viewModels = sections[section].viewModels;
 
       // For each item:
-      for (NSInteger i = 0; i < items.count; i++) {
+      for (NSInteger i = 0; i < viewModels.count; i++) {
         NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
-        [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:items[i]];
+        [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:viewModels[i]];
         [self expectNodeBlockMethodForItemAtIndexPath:indexPath];
       }
     }
@@ -172,9 +196,8 @@
   // For each section:
   // Note: Skip fast enumeration for readability.
   for (NSInteger section = 0; section < sections.count; section++) {
-    NSInteger itemCount = sections[section].count;
     OCMExpect([mockDataSource collectionNode:collectionNode numberOfItemsInSection:section])
-    .andReturn(itemCount);
+    .andReturn(sections[section].viewModels.count);
   }
 }
 
@@ -184,15 +207,24 @@
   .andReturn(viewModel);
 }
 
+- (void)expectContextMethodForSection:(NSInteger)section
+{
+  OCMExpect([mockDataSource collectionNode:collectionNode contextForSection:section])
+  .andReturn(sections[section]);
+}
+
 - (void)expectNodeBlockMethodForItemAtIndexPath:(NSIndexPath *)indexPath
 {
   OCMExpect([mockDataSource collectionNode:collectionNode nodeBlockForItemAtIndexPath:indexPath])
   .andReturn((ASCellNodeBlock)^{
     ASCellNode *node = [ASTestCellNode new];
     // Generating multiple partial mocks of the same class is not thread-safe.
+    id mockNode;
     @synchronized (NSNull.null) {
-      return OCMPartialMock(node);
+      mockNode = OCMPartialMock(node);
     }
+    [mockNode setExpectationOrderMatters:YES];
+    return mockNode;
   });
 }
 
@@ -203,14 +235,19 @@
   XCTAssertEqual(collectionNode.numberOfSections, sections.count);
 
   for (NSInteger section = 0; section < sections.count; section++) {
-    NSArray *items = sections[section];
+    ASTestSection *sectionObject = sections[section];
+    NSArray *viewModels = sectionObject.viewModels;
+
+    // Assert section object
+    XCTAssertEqualObjects([collectionNode contextForSection:section], sectionObject);
+    
     // Assert item count
-    XCTAssertEqual([collectionNode numberOfItemsInSection:section], items.count);
+    XCTAssertEqual([collectionNode numberOfItemsInSection:section], viewModels.count);
-    for (NSInteger item = 0; item < items.count; item++) {
+    for (NSInteger item = 0; item < viewModels.count; item++) {
       // Assert view model
       // Could use pointer equality but the error message is less readable.
       NSIndexPath *indexPath = [NSIndexPath indexPathForItem:item inSection:section];
-      id viewModel = sections[indexPath.section][indexPath.item];
+      id viewModel = viewModels[indexPath.item];
       XCTAssertEqualObjects(viewModel, [collectionNode viewModelForItemAtIndexPath:indexPath]);
       ASCellNode *node = [collectionNode nodeForItemAtIndexPath:indexPath];
       XCTAssertEqualObjects(node.viewModel, viewModel);
@@ -224,30 +261,39 @@
  *
  * skippedReloadIndexPaths are the old index paths for nodes that should use -canUpdateToViewModel: instead of being refetched.
  */
-- (void)performUpdateReloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedItems
+- (void)performUpdateReloadingSections:(NSDictionary<NSNumber *, id> *)reloadedSections
-                     reloadMappings:(NSDictionary<NSIndexPath *, NSIndexPath *> *)reloadMappings
+                        reloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedItems
-                     insertingItems:(NSDictionary<NSIndexPath *, id> *)insertedItems
+                        reloadMappings:(NSDictionary<NSIndexPath *, NSIndexPath *> *)reloadMappings
-                      deletingItems:(NSArray<NSIndexPath *> *)deletedItems
+                        insertingItems:(NSDictionary<NSIndexPath *, id> *)insertedItems
-            skippedReloadIndexPaths:(NSArray<NSIndexPath *> *)skippedReloadIndexPaths
+                         deletingItems:(NSArray<NSIndexPath *> *)deletedItems
+               skippedReloadIndexPaths:(NSArray<NSIndexPath *> *)skippedReloadIndexPaths
 {
   [collectionNode performBatchUpdates:^{
     // First update our data source.
     [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
-      sections[key.section][key.item] = obj;
+      sections[key.section].viewModels[key.item] = obj;
+    }];
+    [reloadedSections enumerateKeysAndObjectsUsingBlock:^(NSNumber * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+      sections[key.integerValue] = obj;
     }];
     
     // Deletion paths, sorted descending
     for (NSIndexPath *indexPath in [deletedItems sortedArrayUsingSelector:@selector(compare:)].reverseObjectEnumerator) {
-      [sections[indexPath.section] removeObjectAtIndex:indexPath.item];
+      [sections[indexPath.section].viewModels removeObjectAtIndex:indexPath.item];
     }
     
     // Insertion paths, sorted ascending.
     NSArray *insertionsSortedAcending = [insertedItems.allKeys sortedArrayUsingSelector:@selector(compare:)];
     for (NSIndexPath *indexPath in insertionsSortedAcending) {
-      [sections[indexPath.section] insertObject:insertedItems[indexPath] atIndex:indexPath.item];
+      [sections[indexPath.section].viewModels insertObject:insertedItems[indexPath] atIndex:indexPath.item];
     }
     
     // Then update the collection node.
+    NSMutableIndexSet *reloadedSectionIndexes = [NSMutableIndexSet indexSet];
+    for (NSNumber *i in reloadedSections) {
+      [reloadedSectionIndexes addIndex:i.integerValue];
+    }
+    [collectionNode reloadSections:reloadedSectionIndexes];
     [collectionNode reloadItemsAtIndexPaths:reloadedItems.allKeys];
     [collectionNode deleteItemsAtIndexPaths:deletedItems];
     [collectionNode insertItemsAtIndexPaths:insertedItems.allKeys];
@@ -260,6 +306,17 @@
     // Combine reloads + inserts and expect them to load content for all of them, in ascending order.
     NSMutableDictionary<NSIndexPath *, id> *insertsPlusReloads = [NSMutableDictionary dictionary];
     [insertsPlusReloads addEntriesFromDictionary:insertedItems];
+
+    // Go through reloaded sections and add all their items into `insertsPlusReloads`
+    [reloadedSectionIndexes enumerateIndexesUsingBlock:^(NSUInteger section, BOOL * _Nonnull stop) {
+      [self expectContextMethodForSection:section];
+      NSArray *viewModels = sections[section].viewModels;
+      for (NSInteger i = 0; i < viewModels.count; i++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
+        insertsPlusReloads[indexPath] = viewModels[i];
+      }
+    }];
+
     [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
       insertsPlusReloads[reloadMappings[key]] = obj;
     }];
@@ -280,6 +337,8 @@
 
 @end
 
+#pragma mark - Other Objects
+
 @implementation ASTestCellNode
 
 - (BOOL)canUpdateToViewModel:(id)viewModel
@@ -289,3 +348,17 @@
 }
 
 @end
+
+@implementation ASTestSection
+@synthesize collectionView;
+@synthesize sectionName;
+
+- (instancetype)init
+{
+  if (self = [super init]) {
+    _viewModels = [NSMutableArray array];
+  }
+  return self;
+}
+
+@end

Tests/ASTestCase.h

@@ -12,6 +12,12 @@
 
 #import <XCTest/XCTest.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
 @interface ASTestCase : XCTestCase
 
+@property (class, nonatomic, nullable, readonly) ASTestCase *currentTestCase;
+
 @end
+
+NS_ASSUME_NONNULL_END

Tests/ASTestCase.m

@@ -12,8 +12,22 @@
 
 #import "ASTestCase.h"
 #import <objc/runtime.h>
+#import <AsyncDisplayKit/AsyncDisplayKit.h>
+#import <OCMock/OCMock.h>
+#import "OCMockObject+ASAdditions.h"
 
-@implementation ASTestCase
+static __weak ASTestCase *currentTestCase;
+
+@implementation ASTestCase {
+  ASWeakSet *registeredMockObjects;
+}
+
+- (void)setUp
+{
+  [super setUp];
+  currentTestCase = self;
+  registeredMockObjects = [ASWeakSet new];
+}
 
 - (void)tearDown
 {
@@ -22,12 +36,15 @@
   for (UIWindow *window in [UIApplication sharedApplication].windows) {
     [window resignKeyWindow];
     window.hidden = YES;
+    window.rootViewController = nil;
     for (UIView *view in window.subviews) {
       [view removeFromSuperview];
     }
   }
   
   // Set nil for all our subclasses' ivars. Use setValue:forKey: so memory is managed correctly.
+  // This is important to do _inside_ the test-perform, so that we catch any issues caused by the
+  // deallocation, and so that we're inside the @autoreleasepool for the test invocation.
   Class c = [self class];
   while (c != [ASTestCase class]) {
     unsigned int ivarCount;
@@ -44,7 +61,48 @@
     c = [c superclass];
   }
 
+  for (OCMockObject *mockObject in registeredMockObjects) {
+    OCMVerifyAll(mockObject);
+    [mockObject stopMocking];
+
+    // Invocations retain arguments, which may cause retain cycles.
+    // Manually clear them all out.
+    NSMutableArray *invocations = object_getIvar(mockObject, class_getInstanceVariable(OCMockObject.class, "invocations"));
+    [invocations removeAllObjects];
+  }
+
+  // Go ahead and spin the run loop before finishing, so the system
+  // unregisters/cleans up whatever possible.
+  [NSRunLoop.mainRunLoop runMode:NSDefaultRunLoopMode beforeDate:NSDate.distantFuture];
+  
   [super tearDown];
 }
 
+- (void)invokeTest
+{
+  // This will call setup, run, then teardown.
+  @autoreleasepool {
+    [super invokeTest];
+  }
+
+  // Now that the autorelease pool is drained, drain the dealloc queue also.
+  [[ASDeallocQueue sharedDeallocationQueue] test_drain];
+}
+
++ (ASTestCase *)currentTestCase
+{
+  return currentTestCase;
+}
+
+@end
+
+@implementation ASTestCase (OCMockObjectRegistering)
+
+- (void)registerMockObject:(id)mockObject
+{
+  @synchronized (registeredMockObjects) {
+    [registeredMockObjects addObject:mockObject];
+  }
+}
+
 @end

Tests/OCMockObject+ASAdditions.h

@@ -14,6 +14,10 @@
 
 @interface OCMockObject (ASAdditions)
 
+/**
+ * NOTE: All OCMockObjects created during an ASTestCase call OCMVerifyAll during -tearDown.
+ */
+
 /**
  * A method to manually specify which optional protocol methods should return YES
  * from -respondsToSelector:.

Tests/OCMockObject+ASAdditions.m

@@ -10,18 +10,32 @@
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 
-#import <OCMock/OCMockObject.h>
+#import "OCMockObject+ASAdditions.h"
+
+#import <OCMock/OCMock.h>
 #import "ASInternalHelpers.h"
 #import <objc/runtime.h>
+#import "ASTestCase.h"
+
+@interface ASTestCase (OCMockObjectRegistering)
+
+- (void)registerMockObject:(id)mockObject;
+
+@end
 
 @implementation OCMockObject (ASAdditions)
 
 + (void)load
 {
-  // Swap [OCProtocolMockObject respondsToSelector:] with [(self) swizzled_protocolMockRespondsToSelector:]
+  // [OCProtocolMockObject respondsToSelector:] <-> [(self) swizzled_protocolMockRespondsToSelector:]
   Method orig = class_getInstanceMethod(OCMockObject.protocolMockObjectClass, @selector(respondsToSelector:));
   Method new = class_getInstanceMethod(self, @selector(swizzled_protocolMockRespondsToSelector:));
   method_exchangeImplementations(orig, new);
+
+  // init <-> swizzled_init
+  Method origInit = class_getInstanceMethod([OCMockObject class], @selector(init));
+  Method newInit = class_getInstanceMethod(self, @selector(swizzled_init));
+  method_exchangeImplementations(origInit, newInit);
 }
 
 /// Since OCProtocolMockObject is private, use this method to get the class.
@@ -120,4 +134,13 @@
   return [self implementsOptionalProtocolMethod:aSelector];
 }
 
+// Whenever a mock object is initted, register it with the current test case
+// so that it gets verified and its invocations are cleared during -tearDown.
+- (instancetype)swizzled_init
+{
+  [self swizzled_init];
+  [ASTestCase.currentTestCase registerMockObject:self];
+  return self;
+}
+
 @end