Prove the change gets the correct nodes with a pair of tests.

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -596,6 +596,7 @@
 		E5711A2C1C840C81009619D4 /* ASIndexedNodeContext.h in Headers */ = {isa = PBXBuildFile; fileRef = E5711A2A1C840C81009619D4 /* ASIndexedNodeContext.h */; };
 		E5711A2E1C840C96009619D4 /* ASIndexedNodeContext.mm in Sources */ = {isa = PBXBuildFile; fileRef = E5711A2D1C840C96009619D4 /* ASIndexedNodeContext.mm */; };
 		E5711A301C840C96009619D4 /* ASIndexedNodeContext.mm in Sources */ = {isa = PBXBuildFile; fileRef = E5711A2D1C840C96009619D4 /* ASIndexedNodeContext.mm */; };
+		F711994E1D20C21100568860 /* ASDisplayNodeExtrasTests.m in Sources */ = {isa = PBXBuildFile; fileRef = F711994D1D20C21100568860 /* ASDisplayNodeExtrasTests.m */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXContainerItemProxy section */
@@ -965,6 +966,7 @@
 		E5711A2A1C840C81009619D4 /* ASIndexedNodeContext.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASIndexedNodeContext.h; sourceTree = "<group>"; };
 		E5711A2D1C840C96009619D4 /* ASIndexedNodeContext.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASIndexedNodeContext.mm; sourceTree = "<group>"; };
 		EFA731F0396842FF8AB635EE /* libPods-AsyncDisplayKitTests.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-AsyncDisplayKitTests.a"; sourceTree = BUILT_PRODUCTS_DIR; };
+		F711994D1D20C21100568860 /* ASDisplayNodeExtrasTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASDisplayNodeExtrasTests.m; sourceTree = "<group>"; };
 		FB07EABBCF28656C6297BC2D /* Pods-AsyncDisplayKitTests.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-AsyncDisplayKitTests.debug.xcconfig"; path = "Pods/Target Support Files/Pods-AsyncDisplayKitTests/Pods-AsyncDisplayKitTests.debug.xcconfig"; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
@@ -1210,6 +1212,7 @@
 				058D0A36195D057000B7D73C /* ASTextNodeTests.m */,
 				058D0A37195D057000B7D73C /* ASTextNodeWordKernerTests.mm */,
 				AEEC47E31C21D3D200EC1693 /* ASVideoNodeTests.m */,
+				F711994D1D20C21100568860 /* ASDisplayNodeExtrasTests.m */,
 				058D09C6195D04C000B7D73C /* Supporting Files */,
 				052EE06A1A15A0D8002C6279 /* TestResources */,
 				2538B6F21BC5D2A2003CA0B4 /* ASCollectionViewFlowLayoutInspectorTests.m */,
@@ -2169,6 +2172,7 @@
 				9F06E5CD1B4CAF4200F015D8 /* ASCollectionViewTests.m in Sources */,
 				2911485C1A77147A005D0878 /* ASControlNodeTests.m in Sources */,
 				CC3B208E1C3F7D0A00798563 /* ASWeakSetTests.m in Sources */,
+				F711994E1D20C21100568860 /* ASDisplayNodeExtrasTests.m in Sources */,
 				ACF6ED5D1B178DC700DA7C62 /* ASDimensionTests.mm in Sources */,
 				058D0A38195D057000B7D73C /* ASDisplayLayerTests.m in Sources */,
 				2538B6F31BC5D2A2003CA0B4 /* ASCollectionViewFlowLayoutInspectorTests.m in Sources */,

AsyncDisplayKitTests/ASDisplayNodeExtrasTests.m

@@ -0,0 +1,76 @@
+//
+//  ASDisplayNodeExtrasTests.m
+//  AsyncDisplayKit
+//
+//  Created by Kiel Gillard on 27/06/2016.
+//  Copyright © 2016 Facebook. All rights reserved.
+//
+
+#import <XCTest/XCTest.h>
+#import <AsyncDisplayKit/AsyncDisplayKit.h>
+#import <AsyncDisplayKit/ASDisplayNodeExtras.h>
+
+@interface ASDisplayNodeExtrasTests : XCTestCase
+
+@end
+
+@interface TestDisplayNode : ASDisplayNode
+@end
+
+@implementation TestDisplayNode
+@end
+
+@implementation ASDisplayNodeExtrasTests
+
+- (void)testShallowFindSubnodesOfSubclass {
+  ASDisplayNode *supernode = [[ASDisplayNode alloc] initWithLayerBlock:^CALayer * _Nonnull{
+    return [CALayer layer];
+  }];
+  NSUInteger count = 10;
+  NSMutableArray *expected = [[NSMutableArray alloc] initWithCapacity:count];
+  for (NSUInteger nodeIndex = 0; nodeIndex < count; nodeIndex++) {
+    TestDisplayNode *node = [[TestDisplayNode alloc] initWithLayerBlock:^CALayer * _Nonnull{
+      return [CALayer layer];
+    }];
+    [supernode addSubnode:node];
+    [expected addObject:node];
+  }
+  NSArray *found = ASDisplayNodeFindAllSubnodesOfClass(supernode, [TestDisplayNode class]);
+  XCTAssertEqualObjects(found, expected, @"Expecting %lu %@ nodes, found %lu", (unsigned long)count, [TestDisplayNode class], (unsigned long)found.count);
+}
+
+- (void)testDeepFindSubnodesOfSubclass {
+  ASDisplayNode *supernode = [[ASDisplayNode alloc] initWithLayerBlock:^CALayer * _Nonnull{
+    return [CALayer layer];
+  }];
+  
+  const NSUInteger count = 2;
+  const NSUInteger levels = 2;
+  const NSUInteger capacity = [[self class] capacityForCount:count levels:levels];
+  NSMutableArray *expected = [[NSMutableArray alloc] initWithCapacity:capacity];
+  
+  [[self class] addSubnodesToNode:supernode number:count remainingLevels:levels accumulated:expected];
+  
+  NSArray *found = ASDisplayNodeFindAllSubnodesOfClass(supernode, [TestDisplayNode class]);
+  XCTAssertEqualObjects(found, expected, @"Expecting %lu %@ nodes, found %lu", (unsigned long)count, [TestDisplayNode class], (unsigned long)found.count);
+}
+
++ (void)addSubnodesToNode:(ASDisplayNode *)supernode number:(NSUInteger)number remainingLevels:(NSUInteger)level accumulated:(inout NSMutableArray *)expected {
+  if (level == 0) return;
+  for (NSUInteger nodeIndex = 0; nodeIndex < number; nodeIndex++) {
+    TestDisplayNode *node = [[TestDisplayNode alloc] initWithLayerBlock:^CALayer * _Nonnull{
+      return [CALayer layer];
+    }];
+    [supernode addSubnode:node];
+    [expected addObject:node];
+    [self addSubnodesToNode:node number:number remainingLevels:(level - 1) accumulated:expected];
+  }
+}
+
+// Graph theory is failing me atm.
++ (NSUInteger)capacityForCount:(NSUInteger)count levels:(NSUInteger)level {
+  if (level == 0) return 0;
+  return pow(count, level) + [self capacityForCount:count levels:(level - 1)];
+}
+
+@end