From ffc15c259ba4e6deebbe932fc0f82230d182c5a7 Mon Sep 17 00:00:00 2001
From: Adlai Holler <guy.on.internet@icloud.com>
Date: Mon, 19 Dec 2016 13:09:45 -0500
Subject: [PATCH] Ensure that all nodes are deallocated inside the thrash
 testing method to avoid polluting other tests

---
 AsyncDisplayKit/ASDisplayNode.mm              |  2 +
 AsyncDisplayKitTests/ASTableViewThrashTests.m | 51 ++++++++++++++-----
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/AsyncDisplayKit/ASDisplayNode.mm b/AsyncDisplayKit/ASDisplayNode.mm
index 215c4dce8f..54299e1d9c 100644
--- a/AsyncDisplayKit/ASDisplayNode.mm
+++ b/AsyncDisplayKit/ASDisplayNode.mm
@@ -2382,6 +2382,7 @@ - (void)__enterHierarchy
   ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert(!_flags.isEnteringHierarchy, @"Should not cause recursive __enterHierarchy");
   ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
+  ASDisplayNodeLogEvent(self, @"enterHierarchy");
   
   // Profiling has shown that locking this method is beneficial, so each of the property accesses don't have to lock and unlock.
   __instanceLock__.lock();
@@ -2428,6 +2429,7 @@ - (void)__exitHierarchy
   ASDisplayNodeAssertMainThread();
   ASDisplayNodeAssert(!_flags.isExitingHierarchy, @"Should not cause recursive __exitHierarchy");
   ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
+  ASDisplayNodeLogEvent(self, @"exitHierarchy");
   
   // Profiling has shown that locking this method is beneficial, so each of the property accesses don't have to lock and unlock.
   __instanceLock__.lock();
diff --git a/AsyncDisplayKitTests/ASTableViewThrashTests.m b/AsyncDisplayKitTests/ASTableViewThrashTests.m
index 55e64b51db..07f23f5f8e 100644
--- a/AsyncDisplayKitTests/ASTableViewThrashTests.m
+++ b/AsyncDisplayKitTests/ASTableViewThrashTests.m
@@ -175,6 +175,12 @@ @interface ASThrashTestNode: ASCellNode
 
 @implementation ASThrashTestNode
 
+- (CGSize)calculateSizeThatFits:(CGSize)constrainedSize
+{
+  ASDisplayNodeAssertFalse(isinf(constrainedSize.width));
+  return CGSizeMake(constrainedSize.width, 44);
+}
+
 @end
 #endif
 
@@ -188,6 +194,8 @@ @interface ASThrashDataSource: NSObject
 @property (nonatomic, strong, readonly) UIWindow *window;
 @property (nonatomic, strong, readonly) TableView *tableView;
 @property (nonatomic, strong) NSArray <ASThrashTestSection *> *data;
+// Only access on main
+@property (nonatomic, strong) ASWeakSet *allNodes;
 @end
 
 
@@ -199,6 +207,7 @@ - (instancetype)initWithData:(NSArray <ASThrashTestSection *> *)data {
     _data = [[NSArray alloc] initWithArray:data copyItems:YES];
     _window = [[UIWindow alloc] initWithFrame:[UIScreen mainScreen].bounds];
     _tableView = [[TableView alloc] initWithFrame:_window.bounds style:UITableViewStylePlain];
+    _allNodes = [[ASWeakSet alloc] init];
     [_window addSubview:_tableView];
 #if USE_UIKIT_REFERENCE
     _tableView.dataSource = self;
@@ -227,6 +236,17 @@ - (CGFloat)tableView:(UITableView *)tableView heightForHeaderInSection:(NSIntege
   return self.data[section].headerHeight;
 }
 
+/// Object passed into predicate is ignored.
+- (NSPredicate *)predicateForDeallocatedHierarchy
+{
+  ASWeakSet *allNodes = self.allNodes;
+  __weak UIWindow *window = _window;
+  __weak ASTableView *view = _tableView;
+  return [NSPredicate predicateWithBlock:^BOOL(id  _Nullable evaluatedObject, NSDictionary<NSString *,id> * _Nullable bindings) {
+    return window == nil && view == nil && allNodes.isEmpty;
+  }];
+}
+
 #if USE_UIKIT_REFERENCE
 
 - (UITableViewCell *)tableView:(UITableView *)tableView cellForRowAtIndexPath:(NSIndexPath *)indexPath {
@@ -240,13 +260,12 @@ - (CGFloat)tableView:(UITableView *)tableView heightForRowAtIndexPath:(NSIndexPa
 
 #else
 
-- (ASCellNodeBlock)tableView:(ASTableView *)tableView nodeBlockForRowAtIndexPath:(NSIndexPath *)indexPath {
-  ASThrashTestItem *item = self.data[indexPath.section].items[indexPath.item];
-  return ^{
-    ASThrashTestNode *node = [[ASThrashTestNode alloc] init];
-    node.item = item;
-    return node;
-  };
+- (ASCellNode *)tableView:(ASTableView *)tableView nodeForRowAtIndexPath:(NSIndexPath *)indexPath
+{
+  ASThrashTestNode *node = [[ASThrashTestNode alloc] init];
+  node.item = self.data[indexPath.section].items[indexPath.item];;
+  [self.allNodes addObject:node];
+  return node;
 }
 
 #endif
@@ -495,11 +514,17 @@ - (void)DISABLED_testRecordedThrashCase {
 - (void)DISABLED_testThrashingWildly {
   for (NSInteger i = 0; i < kThrashingIterationCount; i++) {
     [self setUp];
-    ASThrashDataSource *ds = [[ASThrashDataSource alloc] initWithData:[ASThrashTestSection sectionsWithCount:kInitialSectionCount]];
-    _update = [[ASThrashUpdate alloc] initWithData:ds.data];
-    
-    [self applyUpdate:_update toDataSource:ds];
-    [self verifyDataSource:ds];
+    @autoreleasepool {
+      NSArray *sections = [ASThrashTestSection sectionsWithCount:kInitialSectionCount];
+      _update = [[ASThrashUpdate alloc] initWithData:sections];
+      ASThrashDataSource *ds = [[ASThrashDataSource alloc] initWithData:sections];
+
+      [self applyUpdate:_update toDataSource:ds];
+      [self verifyDataSource:ds];
+      [self expectationForPredicate:[ds predicateForDeallocatedHierarchy] evaluatedWithObject:(id)kCFNull handler:nil];
+    }
+    [self waitForExpectationsWithTimeout:3 handler:nil];
+
     [self tearDown];
   }
 }
@@ -533,7 +558,7 @@ - (void)applyUpdate:(ASThrashUpdate *)update toDataSource:(ASThrashDataSource *)
     [tableView reloadRowsAtIndexPaths:indexPaths withRowAnimation:UITableViewRowAnimationNone];
   }];
   @try {
-    [tableView endUpdates];
+    [tableView endUpdatesAnimated:NO completion:nil];
 #if !USE_UIKIT_REFERENCE
     [tableView waitUntilAllUpdatesAreCommitted];
 #endif