Add W3C element screenshot endpoint support

Summary:
Sometimes it is useful to get a screenshot of the particular element. This is especially useful for non-portrait mode, where things natively available from XCTest are kind of broken.

Unfortunately, this endpoint only works since Xcode9 SDK and requires `screenshotDataForQuality` selector to be available.
Closes facebookarchive/WebDriverAgent#815

Differential Revision: D6543505

Pulled By: marekcirkos

fbshipit-source-id: 08dc4a556c34e2485c478ab5acb344daef5e73b0

Author: Mykola Mokhnach

WebDriverAgent.xcodeproj/project.pbxproj

@@ -10,6 +10,7 @@
 		711084441DA3AA7500F913D6 /* FBXPath.h in Headers */ = {isa = PBXBuildFile; fileRef = 711084421DA3AA7500F913D6 /* FBXPath.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		711084451DA3AA7500F913D6 /* FBXPath.m in Sources */ = {isa = PBXBuildFile; fileRef = 711084431DA3AA7500F913D6 /* FBXPath.m */; };
 		7119E1EC1E891F8600D0B125 /* FBPickerWheelSelectTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 7119E1EB1E891F8600D0B125 /* FBPickerWheelSelectTests.m */; };
+		7126FF0E1FD99C7E00DEFB38 /* FBElementScreenshotTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 7126FF0D1FD99C7E00DEFB38 /* FBElementScreenshotTests.m */; };
 		712A0C851DA3E459007D02E5 /* FBXPathTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 712A0C841DA3E459007D02E5 /* FBXPathTests.m */; };
 		712A0C871DA3E55D007D02E5 /* FBXPath-Private.h in Headers */ = {isa = PBXBuildFile; fileRef = 712A0C861DA3E55D007D02E5 /* FBXPath-Private.h */; };
 		712A0C8C1DA3F25B007D02E5 /* libxml2.tbd in Frameworks */ = {isa = PBXBuildFile; fileRef = 7174AF031D9D39AF008C8AD5 /* libxml2.tbd */; };
@@ -401,6 +402,7 @@
 		711084421DA3AA7500F913D6 /* FBXPath.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = FBXPath.h; sourceTree = "<group>"; };
 		711084431DA3AA7500F913D6 /* FBXPath.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = FBXPath.m; sourceTree = "<group>"; };
 		7119E1EB1E891F8600D0B125 /* FBPickerWheelSelectTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = FBPickerWheelSelectTests.m; sourceTree = "<group>"; };
+		7126FF0D1FD99C7E00DEFB38 /* FBElementScreenshotTests.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = FBElementScreenshotTests.m; sourceTree = "<group>"; };
 		712A0C841DA3E459007D02E5 /* FBXPathTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = FBXPathTests.m; sourceTree = "<group>"; };
 		712A0C861DA3E55D007D02E5 /* FBXPath-Private.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "FBXPath-Private.h"; sourceTree = "<group>"; };
 		7136A4771E8918E60024FC3D /* XCUIElement+FBPickerWheel.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "XCUIElement+FBPickerWheel.h"; sourceTree = "<group>"; };
@@ -1058,6 +1060,7 @@
 			children = (
 				EE9B76991CF799F400275851 /* FBAlertTests.m */,
 				EE26409C1D0EBA25009BE6B0 /* FBElementAttributeTests.m */,
+				7126FF0D1FD99C7E00DEFB38 /* FBElementScreenshotTests.m */,
 				EE006EAC1EB99B15006900A4 /* FBElementVisibilityTests.m */,
 				EE6A89361D0B35920083E92B /* FBFailureProofTestCaseTests.m */,
 				EE1E06DB1D18090F007CF043 /* FBIntegrationTestCase.h */,
@@ -1801,6 +1804,7 @@
 			files = (
 				EE2202131ECC612200A29571 /* FBIntegrationTestCase.m in Sources */,
 				EE22021E1ECC618900A29571 /* FBTapTest.m in Sources */,
+				7126FF0E1FD99C7E00DEFB38 /* FBElementScreenshotTests.m in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};

WebDriverAgentLib/Categories/XCUIDevice+FBHelpers.m

@@ -12,6 +12,7 @@
 #import <arpa/inet.h>
 #import <ifaddrs.h>
 #include <notify.h>
+#import <objc/runtime.h>
 
 #import "FBSpringboardApplication.h"
 #import "FBErrorBuilder.h"
@@ -43,7 +44,7 @@ - (BOOL)fb_goToHomescreenWithError:(NSError **)error
 
 - (NSData *)fb_screenshotWithError:(NSError*__autoreleasing*)error
 {
-  Class xcScreenClass = NSClassFromString(@"XCUIScreen");
+  Class xcScreenClass = objc_lookUpClass("XCUIScreen");
   if (nil == xcScreenClass) {
     NSData *result = [[XCAXClient_iOS sharedClient] screenshotData];
     if (nil == result) {

WebDriverAgentLib/Categories/XCUIElement+FBUtilities.h

@@ -59,6 +59,13 @@ NS_ASSUME_NONNULL_BEGIN
 */
 - (BOOL)fb_waitUntilSnapshotIsStable;
 
+/**
+ Returns screenshot of the particular element
+ @param error If there is an error, upon return contains an NSError object that describes the problem.
+ @return Element screenshot as PNG-encoded data or nil in case of failure
+ */
+- (nullable NSData *)fb_screenshotWithError:(NSError **)error;
+
 @end
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Categories/XCUIElement+FBUtilities.m

@@ -9,6 +9,8 @@
 
 #import "XCUIElement+FBUtilities.h"
 
+#import <objc/runtime.h>
+
 #import "FBAlert.h"
 #import "FBLogger.h"
 #import "FBMacros.h"
@@ -18,6 +20,7 @@
 #import "FBXCodeCompatibility.h"
 #import "XCAXClient_iOS.h"
 #import "XCUIElement+FBWebDriverAttributes.h"
+#import "XCUIScreen.h"
 
 @implementation XCUIElement (FBUtilities)
 
@@ -124,4 +127,48 @@ - (BOOL)fb_waitUntilSnapshotIsStable
   return result;
 }
 
+- (NSData *)fb_screenshotWithError:(NSError **)error
+{
+  if (CGRectIsEmpty(self.frame)) {
+    if (error) {
+      *error = [[FBErrorBuilder.builder withDescription:@"Cannot get a screenshot of zero-sized element"] build];
+    }
+    return nil;
+  }
+
+  Class xcScreenClass = objc_lookUpClass("XCUIScreen");
+  if (nil == xcScreenClass) {
+    if (error) {
+      *error = [[FBErrorBuilder.builder withDescription:@"Element screenshots are only available since Xcode9 SDK"] build];
+    }
+    return nil;
+  }
+
+  XCUIScreen *mainScreen = (XCUIScreen *)[xcScreenClass mainScreen];
+  NSData *result = [mainScreen screenshotDataForQuality:1 rect:self.frame error:error];
+  if (nil == result) {
+    return nil;
+  }
+
+  UIImage *image = [UIImage imageWithData:result];
+  UIInterfaceOrientation orientation = self.application.interfaceOrientation;
+  UIImageOrientation imageOrientation = UIImageOrientationUp;
+  // The received element screenshot will be rotated, if the current interface orientation differs from portrait, so we need to fix that first
+  if (orientation == UIInterfaceOrientationLandscapeRight) {
+    imageOrientation = UIImageOrientationLeft;
+  } else if (orientation == UIInterfaceOrientationLandscapeLeft) {
+    imageOrientation = UIImageOrientationRight;
+  } else if (orientation == UIInterfaceOrientationPortraitUpsideDown) {
+    imageOrientation = UIImageOrientationDown;
+  }
+  CGSize size = image.size;
+  UIGraphicsBeginImageContext(CGSizeMake(size.width, size.height));
+  [[UIImage imageWithCGImage:(CGImageRef)[image CGImage] scale:1.0 orientation:imageOrientation] drawInRect:CGRectMake(0, 0, size.width, size.height)];
+  UIImage *fixedImage = UIGraphicsGetImageFromCurrentImageContext();
+  UIGraphicsEndImageContext();
+
+  // The resulting data is a JPEG image, so we need to convert it to PNG representation
+  return (NSData *)UIImagePNGRepresentation(fixedImage);
+}
+
 @end

WebDriverAgentLib/Commands/FBElementCommands.m

@@ -57,6 +57,7 @@ + (NSArray *)routes
     [[FBRoute POST:@"/element/:uuid/value"] respondWithTarget:self action:@selector(handleSetValue:)],
     [[FBRoute POST:@"/element/:uuid/click"] respondWithTarget:self action:@selector(handleClick:)],
     [[FBRoute POST:@"/element/:uuid/clear"] respondWithTarget:self action:@selector(handleClear:)],
+    [[FBRoute GET:@"/element/:uuid/screenshot"] respondWithTarget:self action:@selector(handleElementScreenshot:)],
     [[FBRoute GET:@"/wda/element/:uuid/accessible"] respondWithTarget:self action:@selector(handleGetAccessible:)],
     [[FBRoute GET:@"/wda/element/:uuid/accessibilityContainer"] respondWithTarget:self action:@selector(handleGetIsAccessibilityContainer:)],
     [[FBRoute POST:@"/wda/element/:uuid/swipe"] respondWithTarget:self action:@selector(handleSwipe:)],
@@ -381,6 +382,19 @@ + (NSArray *)routes
   });
 }
 
++ (id<FBResponsePayload>)handleElementScreenshot:(FBRouteRequest *)request
+{
+  FBElementCache *elementCache = request.session.elementCache;
+  XCUIElement *element = [elementCache elementForUUID:request.parameters[@"uuid"]];
+  NSError *error;
+  NSData *screenshotData = [element fb_screenshotWithError:&error];
+  if (nil == screenshotData) {
+    return FBResponseWithError(error);
+  }
+  NSString *screenshot = [screenshotData base64EncodedStringWithOptions:NSDataBase64Encoding64CharacterLineLength];
+  return FBResponseWithObject(screenshot);
+}
+
 static const CGFloat DEFAULT_OFFSET = (CGFloat)0.2;
 
 + (id<FBResponsePayload>)handleWheelSelect:(FBRouteRequest *)request

WebDriverAgentLib/Utilities/XCTestPrivateSymbols.m

@@ -9,6 +9,8 @@
 
 #import "XCTestPrivateSymbols.h"
 
+#import <objc/runtime.h>
+
 #import "FBRuntimeUtils.h"
 
 NSNumber *FB_XCAXAIsVisibleAttribute;
@@ -38,7 +40,7 @@
 
 void *FBRetrieveXCTestSymbol(const char *name)
 {
-  Class XCTestClass = NSClassFromString(@"XCTestCase");
+  Class XCTestClass = objc_lookUpClass("XCTestCase");
   NSCAssert(XCTestClass != nil, @"XCTest should be already linked", XCTestClass);
   NSString *XCTestBinary = [NSBundle bundleForClass:XCTestClass].executablePath;
   const char *binaryPath = XCTestBinary.UTF8String;

WebDriverAgentTests/IntegrationTests/FBElementScreenshotTests.m

@@ -0,0 +1,47 @@
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+#import <XCTest/XCTest.h>
+
+#import "FBIntegrationTestCase.h"
+#import "XCUIDevice+FBRotation.h"
+#import "XCUIElement+FBUtilities.h"
+
+@interface FBElementScreenshotTests : FBIntegrationTestCase
+@end
+
+@implementation FBElementScreenshotTests
+
+- (void)setUp
+{
+  [super setUp];
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    [self launchApplication];
+    [self goToAlertsPage];
+  });
+}
+
+- (void)testElementScreenshot
+{
+  [[XCUIDevice sharedDevice] fb_setDeviceInterfaceOrientation:UIDeviceOrientationLandscapeLeft];
+  XCUIElement *button = self.testedApplication.buttons[FBShowAlertButtonName];
+  NSError *error = nil;
+  NSData *screenshotData = [button fb_screenshotWithError:&error];
+  if (nil == screenshotData && [error.description containsString:@"available since Xcode9"]) {
+    return;
+  }
+  XCTAssertNotNil(screenshotData);
+  XCTAssertNil(error);
+  UIImage *image = [UIImage imageWithData:screenshotData];
+  XCTAssertNotNil(image);
+  XCTAssertTrue(image.size.width > image.size.height);
+}
+
+@end