From 7836b1c4541c858a062e08caf000218eaeb73e79 Mon Sep 17 00:00:00 2001
From: Oleg Montak <oleg.montak@gmail.com>
Date: Sun, 14 Aug 2022 00:02:28 +0300
Subject: [PATCH] Support viewForRow in grey_pickerColumnSetToValue matcher

---
 AppFramework/Matcher/GREYMatchers.m          | 14 ++++++++++++++
 Tests/TestRig/Sources/PickerViewController.m |  7 +++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/AppFramework/Matcher/GREYMatchers.m b/AppFramework/Matcher/GREYMatchers.m
index 8f3e3009b..3ea5cec78 100644
--- a/AppFramework/Matcher/GREYMatchers.m
+++ b/AppFramework/Matcher/GREYMatchers.m
@@ -23,6 +23,7 @@
 #include <tgmath.h>
 
 #import "UISwitch+GREYApp.h"
+#import "UIView+GREYApp.h"
 #import "GREYElementFinder.h"
 #import "GREYAllOf+Private.h"
 #import "GREYAllOf.h"
@@ -607,6 +608,7 @@ + (void)initialize {
     id<UIPickerViewDelegate> delegate = element.delegate;
     SEL attributedTitleSelector = @selector(pickerView:attributedTitleForRow:forComponent:);
     SEL nonAttributedTitleSelector = @selector(pickerView:titleForRow:forComponent:);
+    SEL viewForRowSelector = @selector(pickerView:viewForRow:forComponent:reusingView:);
     if ([delegate respondsToSelector:attributedTitleSelector]) {
       attributedRowLabel = [delegate pickerView:element
                           attributedTitleForRow:row
@@ -616,6 +618,18 @@ + (void)initialize {
       }
     } else if ([delegate respondsToSelector:nonAttributedTitleSelector]) {
       rowLabel = [delegate pickerView:element titleForRow:row forComponent:column];
+    } else if ([delegate respondsToSelector:viewForRowSelector]) {
+      UIView *rowView = [delegate pickerView:element
+                                  viewForRow:row
+                                  forComponent:column
+                                  reusingView:nil];
+      if (![rowView isKindOfClass:[UILabel class]]) {
+        NSArray<UILabel *> *labels = [rowView grey_childrenAssignableFromClass:[UILabel class]];
+        UILabel *label = (labels.count > 0 ? labels[0] : nil);
+        rowLabel = label.text;
+      } else {
+        rowLabel = [((UILabel *)rowView) text];
+      }
     }
     return rowLabel == value || [rowLabel isEqualToString:value] ||
            [attributedRowLabel.string isEqualToString:value];
diff --git a/Tests/TestRig/Sources/PickerViewController.m b/Tests/TestRig/Sources/PickerViewController.m
index 866b7fe64..3dcbcf61b 100644
--- a/Tests/TestRig/Sources/PickerViewController.m
+++ b/Tests/TestRig/Sources/PickerViewController.m
@@ -58,15 +58,14 @@ - (UIView *)pickerView:(UIPickerView *)pickerView
   UILabel *columnView =
       [[UILabel alloc] initWithFrame:CGRectMake(35, 0, pickerView.frame.size.width / 2,
                                                 pickerView.frame.size.height)];
-  columnView.text = [self pickerView:pickerView titleForRow:row forComponent:component];
+  columnView.text = [self titleForRow:row forComponent:component];
   columnView.textAlignment = NSTextAlignmentCenter;
 
   return columnView;
 }
 
-- (NSString *)pickerView:(UIPickerView *)pickerView
-             titleForRow:(NSInteger)row
-            forComponent:(NSInteger)component {
+- (NSString *)titleForRow:(NSInteger)row
+             forComponent:(NSInteger)component {
   switch (component) {
     case 0:
       return self.customColumn1Array[(NSUInteger)row];