From 5b768aed4729a23189604d40c6d49ff2b6458312 Mon Sep 17 00:00:00 2001
From: dostrander <derko@squareup.com>
Date: Thu, 4 Nov 2021 11:30:28 -0400
Subject: [PATCH] Mimic the logic to find elements in UITableView to
 UICollectionview This uses scrolling to positions to find elements in the
 collectionview rather than using the datasource directly to get the cell.

---
 Sources/KIF/Additions/UIView-KIFAdditions.m | 53 ++++++++++-----------
 1 file changed, 26 insertions(+), 27 deletions(-)

diff --git a/Sources/KIF/Additions/UIView-KIFAdditions.m b/Sources/KIF/Additions/UIView-KIFAdditions.m
index a2776d72..449b8c62 100644
--- a/Sources/KIF/Additions/UIView-KIFAdditions.m
+++ b/Sources/KIF/Additions/UIView-KIFAdditions.m
@@ -303,48 +303,47 @@ - (UIAccessibilityElement *)accessibilityElementMatchingBlock:(BOOL(^)(UIAccessi
         } else if ([self isKindOfClass:[UICollectionView class]]) {
             UICollectionView *collectionView = (UICollectionView *)self;
             
-            NSArray *indexPathsForVisibleItems = [collectionView indexPathsForVisibleItems];
-            
+            NSMutableArray *indexPathsForVisibleItems = [[NSMutableArray alloc] init];
+            [[collectionView visibleCells] enumerateObjectsUsingBlock:^(UICollectionViewCell *cell, NSUInteger idx, BOOL *stop) {
+                NSIndexPath *indexPath = [collectionView indexPathForCell:cell];
+                if (indexPath) {
+                    [indexPathsForVisibleItems addObject:indexPath];
+                }
+            }];
+
+            CFTimeInterval delay = 0.05;
             for (NSUInteger section = 0, numberOfSections = [collectionView numberOfSections]; section < numberOfSections; section++) {
-                for (NSUInteger item = 0, numberOfItems = [collectionView numberOfItemsInSection:section]; item < numberOfItems; item++) {
+                for (NSUInteger row = 0, numberOfItems = [collectionView numberOfItemsInSection:section]; row < numberOfItems; row++) {
                     if (!self.window) {
                         break;
                     }
 
-                    // Skip visible items because they are already handled
-                    NSIndexPath *indexPath = [NSIndexPath indexPathForItem:item inSection:section];
+                    // Skip visible rows because they are already handled.
+                    NSIndexPath *indexPath = [NSIndexPath indexPathForItem:row inSection:section];
                     if ([indexPathsForVisibleItems containsObject:indexPath]) {
-                        continue;
+                        @autoreleasepool {
+                            //scroll to the last row of each section before continuing. Attemps to ensure we can get to sections that are off screen. KIF tests (e.g. testButtonAbsentAfterRemoveFromSuperview) fails without this line. Also without this... we can't expose the next section (in code downstream)
+                            [collectionView scrollToItemAtIndexPath:[indexPathsForVisibleItems lastObject] atScrollPosition:UICollectionViewScrollPositionNone animated:NO];
+                            continue;
+                        }
                     }
-                    
+
                     @autoreleasepool {
-                        // Get the cell directly from the dataSource because UICollectionView will only vend visible cells
-                        UICollectionViewCell *cell = [collectionView.dataSource collectionView:collectionView cellForItemAtIndexPath:indexPath];
+                        // Scroll to the cell and wait for the animation to complete. Using animations here may not be optimal.
+                        CGRect sectionRect = [collectionView layoutAttributesForItemAtIndexPath:indexPath].frame;
+                        [collectionView scrollRectToVisible:sectionRect animated:NO];
 
-                        // The cell contents might change just prior to being displayed, so we simulate the cell appearing onscreen
-                        if ([collectionView.delegate respondsToSelector:@selector(collectionView:willDisplayCell:forItemAtIndexPath:)]) {
-                            [collectionView.delegate collectionView:collectionView willDisplayCell:cell forItemAtIndexPath:indexPath];
-                        }
-                        
+                        UICollectionViewCell *cell = [collectionView cellForItemAtIndexPath:indexPath];
                         UIAccessibilityElement *element = [cell accessibilityElementMatchingBlock:matchBlock notHidden:NO];
-                        
-                        // Remove the cell from the collection view so that it doesn't stick around
-                        [cell removeFromSuperview];
-                        
+
                         // Skip this cell if it isn't the one we're looking for
-                        // Sometimes we get cells with no size here which can cause an endless loop, so we ignore those
-                        if (!element || CGSizeEqualToSize(cell.frame.size, CGSizeZero)) {
+                        if (!element) {
                             continue;
                         }
                     }
-                    
-                    // Scroll to the cell and wait for the animation to complete
-                    CGRect frame = [collectionView.collectionViewLayout layoutAttributesForItemAtIndexPath:indexPath].frame;
-                    [collectionView scrollRectToVisible:frame animated:YES];
+
                     // Note: using KIFRunLoopRunInModeRelativeToAnimationSpeed here may cause tests to stall
-                    CFRunLoopRunInMode(UIApplicationCurrentRunMode, 0.5, false);
-                    
-                    // Now try finding the element again
+                    CFRunLoopRunInMode(UIApplicationCurrentRunMode, delay, false);
                     return [self accessibilityElementMatchingBlock:matchBlock];
                 }
             }