feat: Add a mobile endpoint to retrieve screenshots for each macOS di…

…splay (#120)
appium · Mar 3, 2022 · c1963d9 · c1963d9
1 parent bbd1c66
commit c1963d9
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -490,6 +490,24 @@ formFields | Map or `Array<Pair>` | no | Additional form fields for multipart ht
 
 Base64-encoded content of the recorded media file if `remotePath` parameter is falsy or an empty string.
 
+### macos: screenshots
+
+Retrieves a screenshot of each display available to macOS.
+
+#### Arguments
+
+Name | Type | Required | Description | Example
+--- | --- | --- | --- | ---
+displayId | number | no | Display identifier to take a screenshot for. If not provided then all display screenshots are going to be returned. If no matches were found then an error is thrown. | 1
+
+#### Returns
+
+A dictionary where each key contains a unique display identifier
+and values are dictionaries with following items:
+- `id`: Display identifier
+- `isMain`: Whether this display is the main one
+- `payload`: The actual PNG screenshot data encoded to base64 string
+
 
 ## Application Under Test Concept
 

diff --git a/WebDriverAgentMac/WebDriverAgentLib/Commands/FBScreenshotCommands.m b/WebDriverAgentMac/WebDriverAgentLib/Commands/FBScreenshotCommands.m
@@ -10,6 +10,7 @@
 #import "FBScreenshotCommands.h"
 
 #import "XCTest/XCTest.h"
+#import "FBRouteRequest.h"
 
 @implementation FBScreenshotCommands
 
@@ -21,6 +22,9 @@ + (NSArray *)routes
   @[
     [[FBRoute GET:@"/screenshot"].withoutSession respondWithTarget:self action:@selector(handleGetScreenshot:)],
     [[FBRoute GET:@"/screenshot"] respondWithTarget:self action:@selector(handleGetScreenshot:)],
+
+    [[FBRoute POST:@"/wda/screenshots"].withoutSession respondWithTarget:self action:@selector(handleGetScreenshots:)],
+    [[FBRoute POST:@"/wda/screenshots"] respondWithTarget:self action:@selector(handleGetScreenshots:)],
   ];
 }
 
@@ -39,4 +43,32 @@ + (NSArray *)routes
   return FBResponseWithObject(screenshot);
 }
 
++ (id<FBResponsePayload>)handleGetScreenshots:(FBRouteRequest *)request
+{
+  NSNumber *desiredId = request.arguments[@"displayId"];
+  NSMutableDictionary <NSString *, NSDictionary<NSString *, id> *> *result = [NSMutableDictionary new];
+  NSMutableArray <NSNumber *> *availableDisplayIds = [NSMutableArray new];
+  for (XCUIScreen *screen in XCUIScreen.screens) {
+    NSNumber *displayId = [screen valueForKey:@"_displayID"];
+    if (nil == displayId || (nil != desiredId && ![desiredId isEqualToNumber:displayId])) {
+      continue;
+    }
+
+    [availableDisplayIds addObject:displayId];
+    result[displayId.stringValue] = @{
+      @"id": displayId,
+      @"isMain": [screen valueForKey:@"_isMainScreen"] ?: NSNull.null,
+      @"payload": [screen.screenshot.PNGRepresentation
+                   base64EncodedStringWithOptions:NSDataBase64Encoding64CharacterLineLength] ?: NSNull.null
+    };
+  }
+  if (nil != desiredId && 0 == [result count]) {
+    NSString *message = [NSString stringWithFormat:@"The screen identified by %@ is not available to XCTest. Only the following identifiers are available: %@",
+                         desiredId, availableDisplayIds];
+    return FBResponseWithStatus([FBCommandStatus unableToCaptureScreenErrorWithMessage:message
+                                                                             traceback:nil]);
+  }
+  return FBResponseWithObject(result.copy);
+}
+
 @end
diff --git a/lib/commands/execute.js b/lib/commands/execute.js
@@ -33,6 +33,8 @@ const EXTENSION_COMMANDS_MAPPING = {
 
   startRecordingScreen: 'startRecordingScreen',
   stopRecordingScreen: 'stopRecordingScreen',
+
+  screenshots: 'macosScreenshots',
 };
 
 commands.execute = async function execute (script, args) {

diff --git a/lib/commands/index.js b/lib/commands/index.js
@@ -5,6 +5,7 @@ import sourceCmds from './source';
 import appManagementCmds from './app-management';
 import appleScriptCmds from './applescript';
 import screenRecordingCmds from './record-screen';
+import screenshotsCmds from './screenshots';
 
 const commands = {};
 Object.assign(
@@ -16,6 +17,7 @@ Object.assign(
   appManagementCmds,
   appleScriptCmds,
   screenRecordingCmds,
+  screenshotsCmds,
   // add other command types here
 );
 

diff --git a/lib/commands/screenshots.js b/lib/commands/screenshots.js
@@ -0,0 +1,32 @@
+const commands = {};
+
+/**
+ * @typedef {Object} ScreenshotsInfo
+ *
+ * A dictionary where each key contains a unique display identifier
+ * and values are dictionaries with following items:
+ * - id: Display identifier
+ * - isMain: Whether this display is the main one
+ * - payload: The actual PNG screenshot data encoded to base64 string
+ */
+
+/**
+ * @typedef {Object} ScreenshotsOpts
+ * @property {number?} displayId macOS display identifier to take a screenshot for.
+ * If not provided then screenshots of all displays are going to be returned.
+ * If no matches were found then an error is thrown.
+ */
+
+/**
+ * Retrieves screenshots of each display available to macOS
+ *
+ * @param {ScreenshotsOpts} opts
+ * @returns {ScreenshotsInfo}
+ */
+commands.macosScreenshots = async function macosScreenshots (opts = {}) {
+  const {displayId} = opts;
+  return await this.wda.proxy.command('/wda/screenshots', 'POST', {displayId});
+};
+
+export { commands };
+export default commands;