feat: add currentDisplayId setting for screenshots

WebDriverAgentLib/Utilities/FBConfiguration.h

@@ -368,6 +368,12 @@ typedef NS_ENUM(NSInteger, FBConfigurationKeyboardPreference) {
  */
 @property (atomic, assign) BOOL enforceCustomSnapshots;
 
+/**
+ * The identifier of the display targeted by screenshots.
+ * nil means the main display. Reset to nil for each new session.
+ */
+@property (atomic, copy, nullable) NSNumber *currentDisplayId;
+
 @end
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Utilities/FBConfiguration.m

@@ -357,6 +357,7 @@ - (void)resetSessionSettings
   FBSetCustomParameterForElementSnapshot(FBSnapshotMaxChildrenKey, @INT_MAX);
   self.useClearTextShortcut = YES;
   self.limitXpathContextScope = YES;
+  self.currentDisplayId = nil;
 #if !TARGET_OS_TV && !TARGET_OS_WATCH
   self.screenshotOrientationStorage = UIInterfaceOrientationUnknown;
 #endif

WebDriverAgentLib/Utilities/FBScreen.h

@@ -17,6 +17,23 @@ NS_ASSUME_NONNULL_BEGIN
  */
 + (nullable NSArray<NSDictionary<NSString *, id> *> *)screensWithError:(NSError **)error;
 
+/**
+ The display with the given identifier
+
+ @param displayID The display identifier, as returned by screensWithError:
+ @param error Set if the display is not available
+ @return The matching display or nil
+ */
++ (nullable XCUIScreen *)screenWithDisplayID:(long long)displayID error:(NSError **)error;
+
+/**
+ The display selected by the currentDisplayId setting, or the main display if the setting is not set
+
+ @param error Set if the selected display is no longer available
+ @return The selected display or nil
+ */
++ (nullable XCUIScreen *)currentScreenWithError:(NSError **)error;
+
 /**
  The identifier of the main device's display
  */

WebDriverAgentLib/Utilities/FBScreen.m

@@ -7,6 +7,8 @@
  */
 
 #import "FBScreen.h"
+#import "FBConfiguration.h"
+#import "FBErrorBuilder.h"
 #import "XCUIElement+FBIsVisible.h"
 #import "FBXCodeCompatibility.h"
 #import "XCUIDevice.h"
@@ -40,6 +42,31 @@ @implementation FBScreen
   return result.copy;
 }
 
++ (nullable XCUIScreen *)screenWithDisplayID:(long long)displayID error:(NSError **)error
+{
+  NSArray<XCUIScreen *> *screens = [XCUIDevice.sharedDevice screensOrError:error];
+  if (nil == screens) {
+    return nil;
+  }
+  for (XCUIScreen *screen in screens) {
+    if (screen.displayID == displayID) {
+      return screen;
+    }
+  }
+  [[FBErrorBuilder.builder withDescriptionFormat:@"No display with id %lld is available. Call /wda/screens to list the available displays", displayID]
+   buildError:error];
+  return nil;
+}
+
++ (nullable XCUIScreen *)currentScreenWithError:(NSError **)error
+{
+  NSNumber *displayID = FBConfiguration.sharedInstance.currentDisplayId;
+  if (nil == displayID) {
+    return XCUIScreen.mainScreen;
+  }
+  return [self screenWithDisplayID:displayID.longLongValue error:error];
+}
+
 + (long long)displayID
 {
   return XCUIScreen.mainScreen.displayID;

WebDriverAgentLib/Utilities/FBScreenshot.h

@@ -14,7 +14,7 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FBScreenshot : NSObject
 
 /**
- Retrieves non-scaled screenshot of the whole screen
+ Retrieves non-scaled screenshot of the display selected by the currentDisplayId setting
 
  @param quality The number in range 0-3, where 0 is PNG (lossless), 3 is HEIC (lossless), 1- low quality JPEG and 2 - high quality JPEG
  @param error If there is an error, upon return contains an NSError object that describes the problem.
@@ -24,7 +24,7 @@ NS_ASSUME_NONNULL_BEGIN
                                                    error:(NSError **)error;
 
 /**
- Retrieves non-scaled screenshot of the whole screen
+ Retrieves non-scaled screenshot of the display selected by the currentDisplayId setting
 
  @param screenID The screen identifier to take the screenshot from
  @param compressionQuality Normalized screenshot quality value in range 0..1, where 1 is the best quality

WebDriverAgentLib/Utilities/FBScreenshot.m

@@ -15,6 +15,7 @@
 #import "FBImageProcessor.h"
 #import "FBLogger.h"
 #import "FBMacros.h"
+#import "FBScreen.h"
 #import "FBXCodeCompatibility.h"
 #import "FBXCTestDaemonsProxy.h"
 #import "XCTMessagingChannel_RunnerToDaemon-Protocol.h"
@@ -60,8 +61,11 @@ + (UTType *)imageUtiWithQuality:(NSUInteger)quality
 + (NSData *)takeInOriginalResolutionWithQuality:(NSUInteger)quality
                                           error:(NSError **)error
 {
-  XCUIScreen *mainScreen = XCUIScreen.mainScreen;
-  return [self.class takeWithScreenID:mainScreen.displayID
+  XCUIScreen *screen = [FBScreen currentScreenWithError:error];
+  if (nil == screen) {
+    return nil;
+  }
+  return [self.class takeWithScreenID:screen.displayID
                                 scale:SCREENSHOT_SCALE
                    compressionQuality:[self.class compressionQualityWithQuality:quality]
                             sourceUTI:[self.class imageUtiWithQuality:quality]

WebDriverAgentLib/Utilities/FBSettings.h

@@ -46,5 +46,6 @@ extern NSString *const FB_SETTING_INCLUDE_NATIVE_ACCESSIBILITY_ELEMENT_IN_PAGE_S
 extern NSString *const FB_SETTING_INCLUDE_MIN_MAX_VALUE_IN_PAGE_SOURCE;
 extern NSString *const FB_SETTING_INCLUDE_CUSTOM_ACTIONS_IN_PAGE_SOURCE;
 extern NSString *const FB_SETTING_ENFORCE_CUSTOM_SNAPSHOTS;
+extern NSString *const FB_SETTING_CURRENT_DISPLAY_ID;
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Utilities/FBSettings.m

@@ -42,3 +42,4 @@
 NSString* const FB_SETTING_INCLUDE_MIN_MAX_VALUE_IN_PAGE_SOURCE = @"includeMinMaxValueInPageSource";
 NSString* const FB_SETTING_INCLUDE_CUSTOM_ACTIONS_IN_PAGE_SOURCE = @"includeCustomActionsInPageSource";
 NSString* const FB_SETTING_ENFORCE_CUSTOM_SNAPSHOTS = @"enforceCustomSnapshots";
+NSString* const FB_SETTING_CURRENT_DISPLAY_ID = @"currentDisplayId";

WebDriverAgentLib/Utilities/FBSettingsHandler.m

@@ -12,6 +12,7 @@
 #import "FBClassChainQueryParser.h"
 #import "FBCommandStatus.h"
 #import "FBConfiguration.h"
+#import "FBScreen.h"
 #import "FBSession.h"
 #import "FBSettings.h"
 
@@ -33,6 +34,7 @@ static id FBNormalizedSettingValue(id value)
       FB_SETTING_ACCEPT_ALERT_BUTTON_SELECTOR,
       FB_SETTING_DISMISS_ALERT_BUTTON_SELECTOR,
       FB_SETTING_AUTO_CLICK_ALERT_SELECTOR,
+      FB_SETTING_CURRENT_DISPLAY_ID,
       nil];
   });
   return keys;
@@ -190,6 +192,24 @@ @implementation FBSettingsHandler
       FBConfiguration.sharedInstance.limitXpathContextScope = [value boolValue];
       return nil;
     };
+    map[FB_SETTING_CURRENT_DISPLAY_ID] = ^FBCommandStatus *(FBSession *session, id value) {
+      if (nil == value) {
+        FBConfiguration.sharedInstance.currentDisplayId = nil;
+        return nil;
+      }
+      if (![value isKindOfClass:NSNumber.class]
+          || [value compare:@([value longLongValue])] != NSOrderedSame) {
+        return [FBCommandStatus invalidArgumentErrorWithMessage:@"currentDisplayId must be an integer display id returned by /wda/screens"
+                                                      traceback:nil];
+      }
+      NSError *error;
+      if (nil == [FBScreen screenWithDisplayID:[value longLongValue] error:&error]) {
+        return [FBCommandStatus invalidArgumentErrorWithMessage:error.localizedDescription
+                                                      traceback:nil];
+      }
+      FBConfiguration.sharedInstance.currentDisplayId = @([value longLongValue]);
+      return nil;
+    };
 #if !TARGET_OS_TV && !TARGET_OS_WATCH
     map[FB_SETTING_SCREENSHOT_ORIENTATION] = ^FBCommandStatus *(FBSession *session, id value) {
       NSError *error;
@@ -310,6 +330,9 @@ @implementation FBSettingsHandler
     map[FB_SETTING_LIMIT_XPATH_CONTEXT_SCOPE] = ^id(FBSession *session) {
       return @(FBConfiguration.sharedInstance.limitXpathContextScope);
     };
+    map[FB_SETTING_CURRENT_DISPLAY_ID] = ^id(FBSession *session) {
+      return FBConfiguration.sharedInstance.currentDisplayId ?: @([FBScreen displayID]);
+    };
 #if !TARGET_OS_TV && !TARGET_OS_WATCH
     map[FB_SETTING_SCREENSHOT_ORIENTATION] = ^id(FBSession *session) {
       return FBConfiguration.sharedInstance.humanReadableScreenshotOrientation;

WebDriverAgentTests/IntegrationTests/FBScreenTests.m

@@ -9,7 +9,10 @@
 #import <XCTest/XCTest.h>
 
 #import "FBIntegrationTestCase.h"
+#import "FBConfiguration.h"
 #import "FBScreen.h"
+#import "FBScreenshot.h"
+#import "XCUIScreen.h"
 
 @interface FBScreenTests : FBIntegrationTestCase
 @end
@@ -22,6 +25,12 @@ - (void)setUp
   [self launchApplication];
 }
 
+- (void)tearDown
+{
+  FBConfiguration.sharedInstance.currentDisplayId = nil;
+  [super tearDown];
+}
+
 - (void)testDisplayID
 {
   XCTAssertGreaterThanOrEqual([FBScreen displayID], 0LL);
@@ -50,6 +59,57 @@ - (void)testScreens
   XCTAssertNotNil(mainScreen[@"traits"]);
 }
 
+- (void)testScreenWithDisplayID
+{
+  NSError *error = nil;
+  XCUIScreen *screen = [FBScreen screenWithDisplayID:[FBScreen displayID] error:&error];
+  XCTAssertNil(error);
+  XCTAssertEqual(screen.displayID, [FBScreen displayID]);
+
+  XCTAssertNil([FBScreen screenWithDisplayID:[self unknownDisplayID] error:&error]);
+  XCTAssertNotNil(error);
+}
+
+- (void)testCurrentScreenDefaultsToMainScreen
+{
+  NSError *error = nil;
+  XCTAssertTrue([FBScreen currentScreenWithError:&error].isMainScreen);
+  XCTAssertNil(error);
+}
+
+- (void)testCurrentScreenFollowsSetting
+{
+  FBConfiguration.sharedInstance.currentDisplayId = @([FBScreen displayID]);
+  NSError *error = nil;
+  XCTAssertEqual([FBScreen currentScreenWithError:&error].displayID, [FBScreen displayID]);
+  XCTAssertNil(error);
+  XCTAssertNotNil([FBScreenshot takeInOriginalResolutionWithQuality:0 error:&error]);
+  XCTAssertNil(error);
+
+  FBConfiguration.sharedInstance.currentDisplayId = @([self unknownDisplayID]);
+  XCTAssertNil([FBScreen currentScreenWithError:&error]);
+  XCTAssertNotNil(error);
+  error = nil;
+  XCTAssertNil([FBScreenshot takeInOriginalResolutionWithQuality:0 error:&error]);
+  XCTAssertNotNil(error);
+}
+
+- (void)testSessionResetClearsCurrentDisplay
+{
+  FBConfiguration.sharedInstance.currentDisplayId = @([FBScreen displayID]);
+  [FBConfiguration.sharedInstance resetSessionSettings];
+  XCTAssertNil(FBConfiguration.sharedInstance.currentDisplayId);
+}
+
+- (long long)unknownDisplayID
+{
+  long long maxID = 0;
+  for (NSDictionary<NSString *, id> *screen in [FBScreen screensWithError:nil]) {
+    maxID = MAX(maxID, [screen[@"displayId"] longLongValue]);
+  }
+  return maxID + 1;
+}
+
 - (void)testScreenScale
 {
   XCTAssertTrue([FBScreen scale] >= 2);

test/functional/helpers/display-settings.ts

@@ -0,0 +1,48 @@
+import assert from 'node:assert/strict';
+
+/** Exercise the currentDisplayId setting against a live WDA session. */
+export async function assertCurrentDisplayIdSetting(baseUrl: string, sessionId: string): Promise<void> {
+  const screensResponse = await fetch(`${baseUrl}/wda/screens`);
+  assert.equal(screensResponse.status, 200);
+  const {value: screens} = (await screensResponse.json()) as {
+    value: {displayId: number; isMain: boolean}[];
+  };
+  const main = screens.find((screen) => screen.isMain);
+  assert.ok(main, 'expected a main display');
+  const sessionUrl = `${baseUrl}/session/${sessionId}`;
+  const setDisplay = async (value: unknown) =>
+    await fetch(`${sessionUrl}/appium/settings`, {
+      method: 'POST',
+      headers: {'Content-Type': 'application/json'},
+      body: JSON.stringify({settings: {currentDisplayId: value}}),
+    });
+  const getDisplay = async () => {
+    const response = await fetch(`${sessionUrl}/appium/settings`);
+    assert.equal(response.status, 200);
+    return ((await response.json()) as {value: {currentDisplayId: number}}).value.currentDisplayId;
+  };
+  const assertScreenshot = async () => {
+    const response = await fetch(`${sessionUrl}/screenshot`);
+    assert.equal(response.status, 200);
+    const {value} = (await response.json()) as {value: string};
+    assert.ok(Buffer.from(value, 'base64').length > 0, 'expected screenshot data');
+  };
+
+  assert.equal(await getDisplay(), main.displayId);
+  assert.equal((await setDisplay(main.displayId)).status, 200);
+  assert.equal(await getDisplay(), main.displayId);
+  await assertScreenshot();
+
+  const unknownId = Math.max(...screens.map((screen) => screen.displayId)) + 1;
+  for (const value of [unknownId, 1.5, '1']) {
+    const response = await setDisplay(value);
+    assert.equal(response.status, 400, `expected invalid argument for ${JSON.stringify(value)}`);
+    const {value: body} = (await response.json()) as {value: {error: string}};
+    assert.equal(body.error, 'invalid argument');
+  }
+  assert.equal(await getDisplay(), main.displayId);
+
+  assert.equal((await setDisplay(null)).status, 200);
+  assert.equal(await getDisplay(), main.displayId);
+  await assertScreenshot();
+}

test/functional/webdriveragent-e2e.spec.ts

@@ -5,6 +5,7 @@ import {describe, before, afterEach, it} from 'node:test';
 import type {AppleDevice} from '../../lib/types.js';
 import {WebDriverAgent} from '../../lib/webdriveragent.js';
 import {PLATFORM_NAME, PLATFORM_VERSION} from './desired.js';
+import {assertCurrentDisplayIdSetting} from './helpers/display-settings.js';
 import {getTargetDevice} from './helpers/simulator.js';
 
 const WDA_BASE_URL = 'http://localhost:8100';
@@ -59,6 +60,9 @@ async function assertSessionScopedSourceWorks(): Promise<void> {
     assert.equal(sessionSourceResponse.status, 200);
     const sessionSource = (await sessionSourceResponse.json()) as {value: string};
     assert.ok(sessionSource.value.length > 0, 'expected session-scoped /source to return non-empty page source');
+    if (PLATFORM_NAME.toLowerCase() === 'ios') {
+      await assertCurrentDisplayIdSetting(WDA_BASE_URL, wdaSessionId);
+    }
   } finally {
     await fetch(`${WDA_BASE_URL}/session/${wdaSessionId}`, {method: 'DELETE'});
   }