appium-webdriveragent 4.7.0 → 4.8.0

package/WebDriverAgentLib/Categories/XCUIElement+FBScrolling.m

@@ -13,12 +13,11 @@
 #import "FBLogger.h"
 #import "FBMacros.h"
 #import "FBMathUtils.h"
-#import "FBPredicate.h"
 #import "FBXCodeCompatibility.h"
+#import "FBXCElementSnapshotWrapper.h"
+#import "FBXCElementSnapshotWrapper+Helpers.h"
 #import "XCUIApplication+FBTouchAction.h"
 #import "XCUIElement+FBCaching.h"
-#import "XCElementSnapshot+FBHelpers.h"
-#import "XCElementSnapshot.h"
 #import "XCUIApplication.h"
 #import "XCUICoordinate.h"
 #import "XCUICoordinate+FBFix.h"
@@ -34,7 +33,7 @@ const CGFloat FBScrollTouchProportion = 0.75f;
 
 #if !TARGET_OS_TV
 
-@interface XCElementSnapshot (FBScrolling)
+@interface FBXCElementSnapshotWrapper (FBScrolling)
 
 - (void)fb_scrollUpByNormalizedDistance:(CGFloat)distance inApplication:(XCUIApplication *)application;
 - (void)fb_scrollDownByNormalizedDistance:(CGFloat)distance inApplication:(XCUIApplication *)application;
@@ -49,7 +48,7 @@ const CGFloat FBScrollTouchProportion = 0.75f;
 
 - (BOOL)fb_nativeScrollToVisibleWithError:(NSError **)error
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
   return nil != [self _hitPointByAttemptingToScrollToVisibleSnapshot:snapshot
@@ -58,34 +57,38 @@ const CGFloat FBScrollTouchProportion = 0.75f;
 
 - (void)fb_scrollUpByNormalizedDistance:(CGFloat)distance
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  [snapshot fb_scrollUpByNormalizedDistance:distance inApplication:self.application];
+  [[FBXCElementSnapshotWrapper ensureWrapped:snapshot] fb_scrollUpByNormalizedDistance:distance
+                                                                         inApplication:self.application];
 }
 
 - (void)fb_scrollDownByNormalizedDistance:(CGFloat)distance
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  [snapshot fb_scrollDownByNormalizedDistance:distance inApplication:self.application];
+  [[FBXCElementSnapshotWrapper ensureWrapped:snapshot] fb_scrollDownByNormalizedDistance:distance
+                                                                           inApplication:self.application];
 }
 
 - (void)fb_scrollLeftByNormalizedDistance:(CGFloat)distance
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  [snapshot fb_scrollLeftByNormalizedDistance:distance inApplication:self.application];
+  [[FBXCElementSnapshotWrapper ensureWrapped:snapshot] fb_scrollLeftByNormalizedDistance:distance
+                                                                           inApplication:self.application];
 }
 
 - (void)fb_scrollRightByNormalizedDistance:(CGFloat)distance
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  [snapshot fb_scrollRightByNormalizedDistance:distance inApplication:self.application];
+  [[FBXCElementSnapshotWrapper ensureWrapped:snapshot] fb_scrollRightByNormalizedDistance:distance
+                                                                            inApplication:self.application];
 }
 
 - (BOOL)fb_scrollToVisibleWithError:(NSError **)error
@@ -100,9 +103,11 @@ const CGFloat FBScrollTouchProportion = 0.75f;
                                                         error:error];
 }
 
-- (BOOL)fb_scrollToVisibleWithNormalizedScrollDistance:(CGFloat)normalizedScrollDistance scrollDirection:(FBXCUIElementScrollDirection)scrollDirection error:(NSError **)error
+- (BOOL)fb_scrollToVisibleWithNormalizedScrollDistance:(CGFloat)normalizedScrollDistance
+                                       scrollDirection:(FBXCUIElementScrollDirection)scrollDirection
+                                                 error:(NSError **)error
 {
-  XCElementSnapshot *prescrollSnapshot = [self fb_takeSnapshot];
+  FBXCElementSnapshotWrapper *prescrollSnapshot = [FBXCElementSnapshotWrapper ensureWrapped:[self fb_takeSnapshot]];
   if (prescrollSnapshot.isWDVisible) {
     return YES;
   }
@@ -118,23 +123,30 @@ const CGFloat FBScrollTouchProportion = 0.75f;
     ];
   });
 
-  __block NSArray<XCElementSnapshot *> *cellSnapshots, *visibleCellSnapshots;
-  XCElementSnapshot *scrollView = [prescrollSnapshot fb_parentMatchingOneOfTypes:acceptedParents
-      filter:^(XCElementSnapshot *snapshot) {
-
-         if (![snapshot isWDVisible]) {
-           return NO;
-         }
-
-         cellSnapshots = [snapshot fb_descendantsCellSnapshots];
+  __block NSArray<id<FBXCElementSnapshot>> *cellSnapshots;
+  __block NSMutableArray<id<FBXCElementSnapshot>> *visibleCellSnapshots = [NSMutableArray new];
+  id<FBXCElementSnapshot> scrollView = [prescrollSnapshot fb_parentMatchingOneOfTypes:acceptedParents
+      filter:^(id<FBXCElementSnapshot> snapshot) {
+    FBXCElementSnapshotWrapper *wrappedSnapshot = [FBXCElementSnapshotWrapper ensureWrapped:snapshot];
+    
+    if (![wrappedSnapshot isWDVisible]) {
+      return NO;
+    }
 
-         visibleCellSnapshots = [cellSnapshots filteredArrayUsingPredicate:[FBPredicate predicateWithFormat:@"%K == YES", FBStringify(XCUIElement, fb_isVisible)]];
+    cellSnapshots = [wrappedSnapshot fb_descendantsCellSnapshots];
+    
+    for (id<FBXCElementSnapshot> cellSnapshot in cellSnapshots) {
+      FBXCElementSnapshotWrapper *wrappedCellSnapshot = [FBXCElementSnapshotWrapper ensureWrapped:cellSnapshot];
+      if (wrappedCellSnapshot.wdVisible) {
+        [visibleCellSnapshots addObject:cellSnapshot];
+      }
+    }
 
-         if (visibleCellSnapshots.count > 1) {
-           return YES;
-         }
-         return NO;
-      }];
+    if (visibleCellSnapshots.count > 1) {
+      return YES;
+    }
+    return NO;
+  }];
 
   if (scrollView == nil) {
     return
@@ -143,20 +155,21 @@ const CGFloat FBScrollTouchProportion = 0.75f;
      buildError:error];
   }
 
-  XCElementSnapshot *targetCellSnapshot = [prescrollSnapshot fb_parentCellSnapshot];
-  XCElementSnapshot *lastSnapshot = visibleCellSnapshots.lastObject;
+  id<FBXCElementSnapshot> targetCellSnapshot = [prescrollSnapshot fb_parentCellSnapshot];
+  id<FBXCElementSnapshot> lastSnapshot = visibleCellSnapshots.lastObject;
   // Can't just do indexOfObject, because targetCellSnapshot may represent the same object represented by a member of cellSnapshots, yet be a different object
   // than that member. This reflects the fact that targetCellSnapshot came out of self.fb_parentCellSnapshot, not out of cellSnapshots directly.
   // If the result is NSNotFound, we'll just proceed by scrolling downward/rightward, since NSNotFound will always be larger than the current index.
-  NSUInteger targetCellIndex = [cellSnapshots indexOfObjectPassingTest:^BOOL(XCElementSnapshot *_Nonnull obj, NSUInteger idx, BOOL *_Nonnull stop) {
+  NSUInteger targetCellIndex = [cellSnapshots indexOfObjectPassingTest:^BOOL(id<FBXCElementSnapshot> _Nonnull obj,
+                                                                             NSUInteger idx, BOOL *_Nonnull stop) {
     return [obj _matchesElement:targetCellSnapshot];
   }];
   NSUInteger visibleCellIndex = [cellSnapshots indexOfObject:lastSnapshot];
 
   if (scrollDirection == FBXCUIElementScrollDirectionUnknown) {
     // Try to determine the scroll direction by determining the vector between the first and last visible cells
-    XCElementSnapshot *firstVisibleCell = visibleCellSnapshots.firstObject;
-    XCElementSnapshot *lastVisibleCell = visibleCellSnapshots.lastObject;
+    id<FBXCElementSnapshot> firstVisibleCell = visibleCellSnapshots.firstObject;
+    id<FBXCElementSnapshot> lastVisibleCell = visibleCellSnapshots.lastObject;
     CGVector cellGrowthVector = CGVectorMake(firstVisibleCell.frame.origin.x - lastVisibleCell.frame.origin.x,
                                              firstVisibleCell.frame.origin.y - lastVisibleCell.frame.origin.y
                                              );
@@ -169,17 +182,22 @@ const CGFloat FBScrollTouchProportion = 0.75f;
 
   const NSUInteger maxScrollCount = 25;
   NSUInteger scrollCount = 0;
+  FBXCElementSnapshotWrapper *scrollViewWrapped = [FBXCElementSnapshotWrapper ensureWrapped:scrollView];
   // Scrolling till cell is visible and get current value of frames
   while (![self fb_isEquivalentElementSnapshotVisible:prescrollSnapshot] && scrollCount < maxScrollCount) {
     if (targetCellIndex < visibleCellIndex) {
       scrollDirection == FBXCUIElementScrollDirectionVertical ?
-        [scrollView fb_scrollUpByNormalizedDistance:normalizedScrollDistance inApplication:self.application] :
-        [scrollView fb_scrollLeftByNormalizedDistance:normalizedScrollDistance inApplication:self.application];
+        [scrollViewWrapped fb_scrollUpByNormalizedDistance:normalizedScrollDistance
+                                             inApplication:self.application] :
+        [scrollViewWrapped fb_scrollLeftByNormalizedDistance:normalizedScrollDistance
+                                               inApplication:self.application];
     }
     else {
       scrollDirection == FBXCUIElementScrollDirectionVertical ?
-        [scrollView fb_scrollDownByNormalizedDistance:normalizedScrollDistance inApplication:self.application] :
-        [scrollView fb_scrollRightByNormalizedDistance:normalizedScrollDistance inApplication:self.application];
+        [scrollViewWrapped fb_scrollDownByNormalizedDistance:normalizedScrollDistance
+                                               inApplication:self.application] :
+        [scrollViewWrapped fb_scrollRightByNormalizedDistance:normalizedScrollDistance
+                                                inApplication:self.application];
     }
     scrollCount++;
     // Wait for scroll animation
@@ -196,28 +214,32 @@ const CGFloat FBScrollTouchProportion = 0.75f;
   // Cell is now visible, but it might be only partialy visible, scrolling till whole frame is visible.
   // Sometimes, attempting to grab the parent snapshot of the target cell after scrolling is complete causes a stale element reference exception.
   // Trying fb_cachedSnapshot first
-  targetCellSnapshot = [([self fb_cachedSnapshot] ?: [self fb_takeSnapshot]) fb_parentCellSnapshot];
-  CGRect visibleFrame = targetCellSnapshot.fb_visibleFrameWithFallback;
+  FBXCElementSnapshotWrapper *targetCellSnapshotWrapped = [FBXCElementSnapshotWrapper ensureWrapped:([self fb_cachedSnapshot] ?: [self fb_takeSnapshot])];
+  targetCellSnapshot = [targetCellSnapshotWrapped fb_parentCellSnapshot];
+  CGRect visibleFrame = [FBXCElementSnapshotWrapper ensureWrapped:targetCellSnapshot].fb_visibleFrameWithFallback;
   
   CGVector scrollVector = CGVectorMake(visibleFrame.size.width - targetCellSnapshot.frame.size.width,
                                        visibleFrame.size.height - targetCellSnapshot.frame.size.height
                                        );
-  if (![scrollView fb_scrollByVector:scrollVector inApplication:self.application error:error]) {
-    return NO;
-  }
-  return YES;
+  return [scrollViewWrapped fb_scrollByVector:scrollVector
+                                inApplication:self.application
+                                        error:error];
 }
 
-- (BOOL)fb_isEquivalentElementSnapshotVisible:(XCElementSnapshot *)snapshot
+- (BOOL)fb_isEquivalentElementSnapshotVisible:(id<FBXCElementSnapshot>)snapshot
 {
-  if (snapshot.isWDVisible) {
+  FBXCElementSnapshotWrapper *wrappedSnapshot = [FBXCElementSnapshotWrapper ensureWrapped:snapshot];
+  
+  if (wrappedSnapshot.isWDVisible) {
     return YES;
   }
 
-  XCElementSnapshot *appSnapshot = [self.application fb_takeSnapshot];
-  for (XCElementSnapshot *elementSnapshot in appSnapshot._allDescendants.copy) {
+  id<FBXCElementSnapshot> appSnapshot = [self.application fb_takeSnapshot];
+  for (id<FBXCElementSnapshot> elementSnapshot in appSnapshot._allDescendants.copy) {
+    FBXCElementSnapshotWrapper *wrappedElementSnapshot = [FBXCElementSnapshotWrapper ensureWrapped:elementSnapshot];
     // We are comparing pre-scroll snapshot so frames are irrelevant.
-    if ([snapshot fb_framelessFuzzyMatchesElement:elementSnapshot] && elementSnapshot.isWDVisible) {
+    if ([wrappedSnapshot fb_framelessFuzzyMatchesElement:elementSnapshot]
+        && wrappedElementSnapshot.isWDVisible) {
       return YES;
     }
   }
@@ -227,7 +249,7 @@ const CGFloat FBScrollTouchProportion = 0.75f;
 @end
 
 
-@implementation XCElementSnapshot (FBScrolling)
+@implementation FBXCElementSnapshotWrapper (FBScrolling)
 
 - (CGRect)scrollingFrame
 {

package/WebDriverAgentLib/Categories/XCUIElement+FBTyping.m

@@ -13,7 +13,8 @@
 #import "FBErrorBuilder.h"
 #import "FBKeyboard.h"
 #import "NSString+FBVisualLength.h"
-#import "XCElementSnapshot+FBHelpers.h"
+#import "FBXCElementSnapshotWrapper.h"
+#import "FBXCElementSnapshotWrapper+Helpers.h"
 #import "XCUIElement+FBCaching.h"
 #import "XCUIElement+FBTap.h"
 #import "XCUIElement+FBUtilities.h"
@@ -39,19 +40,19 @@
 @end
 
 
-@interface XCElementSnapshot (FBKeyboardFocus)
+@interface FBXCElementSnapshotWrapper (FBKeyboardFocus)
 
 - (BOOL)fb_hasKeyboardFocus;
 
 @end
 
-@implementation XCElementSnapshot (FBKeyboardFocus)
+@implementation FBXCElementSnapshotWrapper (FBKeyboardFocus)
 
 - (BOOL)fb_hasKeyboardFocus
 {
   // https://developer.apple.com/documentation/xctest/xcuielement/1500968-typetext?language=objc
   // > The element or a descendant must have keyboard focus; otherwise an error is raised.
-  return self.hasKeyboardFocus || [self descendantsByFilteringWithBlock:^BOOL(XCElementSnapshot *snapshot) {
+  return self.hasKeyboardFocus || [self descendantsByFilteringWithBlock:^BOOL(id<FBXCElementSnapshot> snapshot) {
     return snapshot.hasKeyboardFocus;
   }].count > 0;
 }
@@ -61,7 +62,7 @@
 
 @implementation XCUIElement (FBTyping)
 
-- (void)fb_prepareForTextInputWithSnapshot:(XCElementSnapshot *)snapshot
+- (void)fb_prepareForTextInputWithSnapshot:(FBXCElementSnapshotWrapper *)snapshot
 {
   if (snapshot.fb_hasKeyboardFocus) {
     return;
@@ -92,10 +93,10 @@
           frequency:(NSUInteger)frequency
               error:(NSError **)error
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  [self fb_prepareForTextInputWithSnapshot:snapshot];
+  [self fb_prepareForTextInputWithSnapshot:[FBXCElementSnapshotWrapper ensureWrapped:snapshot]];
   if (shouldClear && ![self fb_clearTextWithSnapshot:self.lastSnapshot
                                shouldPrepareForInput:NO
                                                error:error]) {
@@ -106,15 +107,15 @@
 
 - (BOOL)fb_clearTextWithError:(NSError **)error
 {
-  XCElementSnapshot *snapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> snapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
-  return [self fb_clearTextWithSnapshot:snapshot
+  return [self fb_clearTextWithSnapshot:[FBXCElementSnapshotWrapper ensureWrapped:snapshot]
                   shouldPrepareForInput:YES
                                   error:error];
 }
 
-- (BOOL)fb_clearTextWithSnapshot:(XCElementSnapshot *)snapshot
+- (BOOL)fb_clearTextWithSnapshot:(FBXCElementSnapshotWrapper *)snapshot
            shouldPrepareForInput:(BOOL)shouldPrepareForInput
                            error:(NSError **)error
 {

package/WebDriverAgentLib/Categories/XCUIElement+FBUID.h

@@ -7,24 +7,29 @@
  * of patent rights can be found in the PATENTS file in the same directory.
  */
 
-#import <WebDriverAgentLib/XCElementSnapshot.h>
-#import <WebDriverAgentLib/XCUIElement.h>
+#import "FBXCElementSnapshotWrapper.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
 @interface XCUIElement (FBUID)
 
-/*! Represents unique internal element identifier, which is the same for an element and its snapshot */
+/*! Represents unique internal element identifier, which is the same for an element and its snapshot as UUIDv4 */
 @property (nonatomic, nullable, readonly, copy) NSString *fb_uid;
 
+/*! Represents unique internal element identifier, which is the same for an element and its snapshot */
+@property (nonatomic, readonly) unsigned long long fb_accessibiltyId;
+
 @end
 
 
-@interface XCElementSnapshot (FBUID)
+@interface FBXCElementSnapshotWrapper (FBUID)
 
-/*! Represents unique internal element identifier, which is the same for an element and its snapshot */
+/*! Represents unique internal element identifier, which is the same for an element and its snapshot as UUIDv4 */
 @property (nonatomic, nullable, readonly, copy) NSString *fb_uid;
 
+/*! Represents unique internal element identifier, which is the same for an element and its snapshot */
+@property (nonatomic, readonly) unsigned long long fb_accessibiltyId;
+
 @end
 
 NS_ASSUME_NONNULL_END

package/WebDriverAgentLib/Categories/XCUIElement+FBUID.m

@@ -15,20 +15,32 @@
 
 @implementation XCUIElement (FBUID)
 
+- (unsigned long long)fb_accessibiltyId
+{
+  return [FBElementUtils idWithAccessibilityElement:([self isKindOfClass:XCUIApplication.class]
+                                                     ? [(XCUIApplication *)self accessibilityElement]
+                                                     : [self fb_takeSnapshot].accessibilityElement)];
+}
+
 - (NSString *)fb_uid
 {
   return [self isKindOfClass:XCUIApplication.class]
     ? [FBElementUtils uidWithAccessibilityElement:[(XCUIApplication *)self accessibilityElement]]
-    : [self fb_takeSnapshot].fb_uid;
+    : [FBXCElementSnapshotWrapper ensureWrapped:[self fb_takeSnapshot]].fb_uid;
 }
 
 @end
 
-@implementation XCElementSnapshot (FBUID)
+@implementation FBXCElementSnapshotWrapper (FBUID)
+
+- (unsigned long long)fb_accessibiltyId
+{
+  return [FBElementUtils idWithAccessibilityElement:self.accessibilityElement];
+}
 
 - (NSString *)fb_uid
 {
-  return [FBElementUtils uidWithAccessibilityElement:self.accessibilityElement];
+  return [FBElementUtils uidWithAccessibilityElement:[self accessibilityElement]];
 }
 
 @end

package/WebDriverAgentLib/Categories/XCUIElement+FBUtilities.h

@@ -8,8 +8,8 @@
  */
 
 #import <XCTest/XCTest.h>
-#import <WebDriverAgentLib/XCElementSnapshot.h>
 #import <WebDriverAgentLib/FBElement.h>
+#import "FBXCElementSnapshot.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -24,7 +24,7 @@ NS_ASSUME_NONNULL_BEGIN
  @return The recent snapshot of the element
  @throws FBStaleElementException if the element is not present in DOM and thus no snapshot could be made
  */
-- (XCElementSnapshot *)fb_takeSnapshot;
+- (id<FBXCElementSnapshot>)fb_takeSnapshot;
 
 /**
  Extracts the cached element snapshot from its query.
@@ -34,7 +34,7 @@ NS_ASSUME_NONNULL_BEGIN
 
  @return Either the cached snapshot or nil
  */
-- (nullable XCElementSnapshot *)fb_cachedSnapshot;
+- (nullable id<FBXCElementSnapshot>)fb_cachedSnapshot;
 
 /**
  Gets the most recent snapshot of the current element and already resolves the accessibility attributes
@@ -50,7 +50,7 @@ NS_ASSUME_NONNULL_BEGIN
  attributes resolved if there was a failure while resolving additional attributes
  @throws FBStaleElementException if the element is not present in DOM and thus no snapshot could be made
  */
-- (nullable XCElementSnapshot *)fb_snapshotWithAllAttributesAndMaxDepth:(nullable NSNumber *)maxDepth;
+- (nullable id<FBXCElementSnapshot>)fb_snapshotWithAllAttributesAndMaxDepth:(nullable NSNumber *)maxDepth;
 
 /**
  Gets the most recent snapshot of the current element with given attributes resolved.
@@ -67,8 +67,8 @@ NS_ASSUME_NONNULL_BEGIN
  attributes resolved if there was a failure while resolving additional attributes
  @throws FBStaleElementException if the element is not present in DOM and thus no snapshot could be made
 */
-- (nullable XCElementSnapshot *)fb_snapshotWithAttributes:(nullable NSArray<NSString *> *)attributeNames
-                                                 maxDepth:(nullable NSNumber *)maxDepth;
+- (nullable id<FBXCElementSnapshot>)fb_snapshotWithAttributes:(nullable NSArray<NSString *> *)attributeNames
+                                                      maxDepth:(nullable NSNumber *)maxDepth;
 
 /**
  Filters elements by matching them to snapshots from the corresponding array
@@ -80,7 +80,7 @@ NS_ASSUME_NONNULL_BEGIN
 
  @return Array of filtered elements, which have matches in snapshots array
  */
-- (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<XCElementSnapshot *> *)snapshots
+- (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<id<FBXCElementSnapshot>> *)snapshots
                                                       selfUID:(nullable NSString *)selfUID
                                                  onlyChildren:(BOOL)onlyChildren;
 

package/WebDriverAgentLib/Categories/XCUIElement+FBUtilities.m

@@ -22,6 +22,8 @@
 #import "FBScreenshot.h"
 #import "FBXCAXClientProxy.h"
 #import "FBXCodeCompatibility.h"
+#import "FBXCElementSnapshot.h"
+#import "FBXCElementSnapshotWrapper+Helpers.h"
 #import "XCUIApplication.h"
 #import "XCUIApplication+FBQuiescence.h"
 #import "XCUIApplicationImpl.h"
@@ -41,7 +43,7 @@
 
 @implementation XCUIElement (FBUtilities)
 
-- (XCElementSnapshot *)fb_takeSnapshot
+- (id<FBXCElementSnapshot>)fb_takeSnapshot
 {
   NSError *error = nil;
   self.fb_isResolvedFromCache = @(NO);
@@ -61,12 +63,12 @@
   return self.lastSnapshot;
 }
 
-- (XCElementSnapshot *)fb_cachedSnapshot
+- (id<FBXCElementSnapshot>)fb_cachedSnapshot
 {
   return [self.query fb_cachedSnapshot];
 }
 
-- (nullable XCElementSnapshot *)fb_snapshotWithAllAttributesAndMaxDepth:(NSNumber *)maxDepth
+- (nullable id<FBXCElementSnapshot>)fb_snapshotWithAllAttributesAndMaxDepth:(NSNumber *)maxDepth
 {
   NSMutableArray *allNames = [NSMutableArray arrayWithArray:FBStandardAttributeNames()];
   [allNames addObjectsFromArray:FBCustomAttributeNames()];
@@ -74,11 +76,11 @@
                                 maxDepth:maxDepth];
 }
 
-- (nullable XCElementSnapshot *)fb_snapshotWithAttributes:(NSArray<NSString *> *)attributeNames
-                                                 maxDepth:(NSNumber *)maxDepth
+- (nullable id<FBXCElementSnapshot>)fb_snapshotWithAttributes:(NSArray<NSString *> *)attributeNames
+                                                     maxDepth:(NSNumber *)maxDepth
 {
   NSSet<NSString *> *standardAttributes = [NSSet setWithArray:FBStandardAttributeNames()];
-  XCElementSnapshot *snapshot = self.fb_takeSnapshot;
+  id<FBXCElementSnapshot> snapshot = self.fb_takeSnapshot;
   NSTimeInterval axTimeout = FBConfiguration.customSnapshotTimeout;
   if (nil == attributeNames
       || [[NSSet setWithArray:attributeNames] isSubsetOfSet:standardAttributes]
@@ -87,7 +89,7 @@
     return snapshot;
   }
 
-  XCAccessibilityElement *axElement = snapshot.accessibilityElement;
+  id<FBXCAccessibilityElement> axElement = snapshot.accessibilityElement;
   if (nil == axElement) {
     return nil;
   }
@@ -101,16 +103,17 @@
   }
 
   NSError *error;
-  XCElementSnapshot *snapshotWithAttributes = [FBXCAXClientProxy.sharedClient snapshotForElement:axElement
-                                                                                      attributes:attributeNames
-                                                                                        maxDepth:maxDepth
-                                                                                           error:&error];
+  id<FBXCElementSnapshot> snapshotWithAttributes = [FBXCAXClientProxy.sharedClient snapshotForElement:axElement
+                                                                                           attributes:attributeNames
+                                                                                             maxDepth:maxDepth
+                                                                                                error:&error];
   if (nil == snapshotWithAttributes) {
+    NSString *description = [FBXCElementSnapshotWrapper ensureWrapped:snapshot].fb_description;
     [FBLogger logFmt:@"Cannot take a snapshot with attribute(s) %@ of '%@' after %.2f seconds",
-     attributeNames, snapshot.fb_description, axTimeout];
+     attributeNames, description, axTimeout];
     [FBLogger logFmt:@"This timeout could be customized via '%@' setting", FB_SETTING_CUSTOM_SNAPSHOT_TIMEOUT];
     [FBLogger logFmt:@"Internal error: %@", error.localizedDescription];
-    [FBLogger logFmt:@"Falling back to the default snapshotting mechanism for the element '%@' (some attribute values, like visibility or accessibility might not be precise though)", snapshot.fb_description];
+    [FBLogger logFmt:@"Falling back to the default snapshotting mechanism for the element '%@' (some attribute values, like visibility or accessibility might not be precise though)", description];
     snapshotWithAttributes = self.lastSnapshot;
   } else {
     self.lastSnapshot = snapshotWithAttributes;
@@ -122,19 +125,22 @@
   return snapshotWithAttributes;
 }
 
-- (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<XCElementSnapshot *> *)snapshots
+- (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<id<FBXCElementSnapshot>> *)snapshots
                                                       selfUID:(NSString *)selfUID
                                                  onlyChildren:(BOOL)onlyChildren
 {
   if (0 == snapshots.count) {
     return @[];
   }
-  NSArray<NSString *> *sortedIds = [snapshots valueForKey:FBStringify(XCUIElement, wdUID)];
+  NSMutableArray<NSString *> *sortedIds = [NSMutableArray new];
+  for (id<FBXCElementSnapshot> snapshot in snapshots) {
+    [sortedIds addObject:(NSString *)[FBXCElementSnapshotWrapper ensureWrapped:snapshot].fb_uid];
+  }
   NSMutableArray<XCUIElement *> *matchedElements = [NSMutableArray array];
   NSString *uid = selfUID;
   if (nil == uid) {
     uid = self.fb_isResolvedFromCache.boolValue
-      ? self.lastSnapshot.fb_uid
+      ? [FBXCElementSnapshotWrapper ensureWrapped:self.lastSnapshot].fb_uid
       : self.fb_uid;
   }
   if ([sortedIds containsObject:uid]) {
@@ -151,7 +157,10 @@
   XCUIElementQuery *query = onlyChildren
     ? [self.fb_query childrenMatchingType:type]
     : [self.fb_query descendantsMatchingType:type];
-  query = [query matchingPredicate:[NSPredicate predicateWithFormat:@"%K IN %@", FBStringify(XCUIElement, wdUID), sortedIds]];
+  NSPredicate *predicate = [NSPredicate predicateWithBlock:^BOOL(id<FBXCElementSnapshot> snapshot, NSDictionary *bindings) {
+    return [sortedIds containsObject:(NSString *)[FBXCElementSnapshotWrapper ensureWrapped:snapshot].fb_uid];
+  }];
+  query = [query matchingPredicate:predicate];
   if (1 == snapshots.count) {
     XCUIElement *result = query.fb_firstMatch;
     result.fb_isResolvedNatively = @NO;
@@ -206,7 +215,7 @@
 
 - (NSData *)fb_screenshotWithError:(NSError **)error
 {
-  XCElementSnapshot *selfSnapshot = self.fb_isResolvedFromCache.boolValue
+  id<FBXCElementSnapshot> selfSnapshot = self.fb_isResolvedFromCache.boolValue
     ? self.lastSnapshot
     : self.fb_takeSnapshot;
   if (CGRectIsEmpty(selfSnapshot.frame)) {
@@ -221,8 +230,8 @@
   UIInterfaceOrientation orientation = self.application.interfaceOrientation;
   if (orientation == UIInterfaceOrientationLandscapeLeft || orientation == UIInterfaceOrientationLandscapeRight) {
     // Workaround XCTest bug when element frame is returned as in portrait mode even if the screenshot is rotated
-    NSArray<XCElementSnapshot *> *ancestors = selfSnapshot.fb_ancestors;
-    XCElementSnapshot *parentWindow = nil;
+    NSArray<id<FBXCElementSnapshot>> *ancestors = [FBXCElementSnapshotWrapper ensureWrapped:selfSnapshot].fb_ancestors;
+    id<FBXCElementSnapshot> parentWindow = nil;
     if (1 == ancestors.count) {
       parentWindow = selfSnapshot;
     } else if (ancestors.count > 1) {

package/WebDriverAgentLib/Categories/XCUIElement+FBWebDriverAttributes.h

@@ -8,8 +8,8 @@
  */
 
 #import <WebDriverAgentLib/FBElement.h>
-#import <WebDriverAgentLib/XCElementSnapshot.h>
 #import <WebDriverAgentLib/XCUIElement.h>
+#import "FBXCElementSnapshotWrapper.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -18,7 +18,7 @@ NS_ASSUME_NONNULL_BEGIN
 @end
 
 
-@interface XCElementSnapshot (WebDriverAttributes) <FBElement>
+@interface FBXCElementSnapshotWrapper (WebDriverAttributes) <FBElement>
 
 @end