fix: detach screen recording promise before XCTest attachment cleanup

WebDriverAgentLib/Routing/FBScreenRecordingContainer.m

@@ -9,14 +9,17 @@
 
 #import "FBScreenRecordingContainer.h"
 
+#import "FBLogger.h"
 #import "FBScreenRecordingPromise.h"
 
 @interface FBScreenRecordingContainer ()
 
 @property (readwrite) NSUInteger fps;
 @property (readwrite) long long codec;
-@property (readwrite) FBScreenRecordingPromise* screenRecordingPromise;
-@property (readwrite) NSNumber *startedAt;
+@property (readwrite, nullable) FBScreenRecordingPromise* screenRecordingPromise;
+@property (readwrite, nullable) NSNumber *startedAt;
+
+- (void)performAttachmentCleanupForPromise:(FBScreenRecordingPromise *)promise;
 
 @end
 
@@ -44,45 +47,72 @@ - (void)storeScreenRecordingPromise:(FBScreenRecordingPromise *)screenRecordingP
   }
 }
 
-- (void)reset;
+- (nullable FBScreenRecordingPromise *)detachScreenRecordingPromiseMatching:(nullable FBScreenRecordingPromise *)expectedPromise
+                                                                   matchAny:(BOOL)matchAny
 {
   @synchronized (self) {
+    if (nil == self.screenRecordingPromise || (!matchAny && self.screenRecordingPromise != expectedPromise)) {
+      return nil;
+    }
+    FBScreenRecordingPromise *detachedPromise = self.screenRecordingPromise;
     self.fps = 0;
     self.codec = 0;
-    if (nil != self.screenRecordingPromise) {
-      [XCTContext runActivityNamed:@"Video Cleanup" block:^(id<XCTActivity> activity){
-        [activity addAttachment:(XCTAttachment *)self.screenRecordingPromise.nativePromise];
-      }];
-      self.screenRecordingPromise = nil;
-    }
+    self.screenRecordingPromise = nil;
     self.startedAt = nil;
+    return detachedPromise;
+  }
+}
+
+- (void)performAttachmentCleanupForPromise:(FBScreenRecordingPromise *)promise
+{
+  [XCTContext runActivityNamed:@"Video Cleanup" block:^(id<XCTActivity> activity){
+    [activity addAttachment:(XCTAttachment *)promise.nativePromise];
+  }];
+}
+
+- (void)cleanUpScreenRecordingPromise:(nullable FBScreenRecordingPromise *)promise
+{
+  if (nil == promise || nil == promise.nativePromise) {
+    return;
   }
+  @try {
+    [self performAttachmentCleanupForPromise:promise];
+  } @catch (NSException *exception) {
+    [FBLogger logFmt:@"Failed to clean up screen recording attachment: %@", exception.reason];
+  }
+}
+
+- (void)reset;
+{
+  FBScreenRecordingPromise *detachedPromise = [self detachScreenRecordingPromiseMatching:nil matchAny:YES];
+  [self cleanUpScreenRecordingPromise:detachedPromise];
 }
 
 - (BOOL)resetIfPromiseIs:(FBScreenRecordingPromise *)screenRecordingPromise
 {
-  // @synchronized is recursive, so -reset may take the very same lock again below.
-  @synchronized (self) {
-    if (self.screenRecordingPromise != screenRecordingPromise) {
-      return NO;
-    }
-    [self reset];
-    return YES;
+  FBScreenRecordingPromise *detachedPromise = [self detachScreenRecordingPromiseMatching:screenRecordingPromise
+                                                                                matchAny:NO];
+  if (nil == detachedPromise) {
+    return NO;
   }
+  [self cleanUpScreenRecordingPromise:detachedPromise];
+  return YES;
 }
 
 - (nullable NSDictionary *)toDictionary
 {
-  if (nil == self.screenRecordingPromise) {
-    return nil;
-  }
+  @synchronized (self) {
+    if (nil == self.screenRecordingPromise) {
+      return nil;
+    }
 
-  return @{
-    @"fps": @(self.fps),
-    @"codec": @(self.codec),
-    @"uuid": [self.screenRecordingPromise identifier].UUIDString ?: [NSNull null],
-    @"startedAt": self.startedAt ?: [NSNull null],
-  };
+    return @{
+      @"fps": @(self.fps),
+      @"codec": @(self.codec),
+      @"uuid": [self.screenRecordingPromise identifier].UUIDString ?: [NSNull null],
+      @"startedAt": self.startedAt ?: [NSNull null],
+    };
+  }
 }
 
 @end

WebDriverAgentTests/UnitTests/FBSessionTests.m

@@ -12,8 +12,45 @@
 #import "FBSession-Private.h"
 #import "FBExceptions.h"
 #import "FBConfiguration.h"
+#import "FBScreenRecordingContainer.h"
+#import "FBScreenRecordingPromise.h"
 #import "XCUIApplicationDouble.h"
 
+@interface FBScreenRecordingContainer (ScreenRecordingContainerTests)
+- (void)performAttachmentCleanupForPromise:(FBScreenRecordingPromise *)promise;
+@end
+
+@interface FBNativeScreenRecordingPromiseDouble : NSObject
+@property (nonatomic, strong) NSUUID *identifier;
+@end
+
+@implementation FBNativeScreenRecordingPromiseDouble
+- (nullable id)valueForKey:(NSString *)key
+{
+  if ([key isEqualToString:@"_UUID"]) {
+    return self.identifier;
+  }
+  return [super valueForKey:key];
+}
+@end
+
+@interface FBTestableScreenRecordingContainer : FBScreenRecordingContainer
+@property (nonatomic, copy, nullable) void (^attachmentCleanupHandler)(FBScreenRecordingContainer *targetContainer, FBScreenRecordingPromise *promise);
+@property (nonatomic) NSUInteger cleanupCallCount;
+@end
+
+@implementation FBTestableScreenRecordingContainer
+- (void)performAttachmentCleanupForPromise:(FBScreenRecordingPromise *)promise
+{
+  self.cleanupCallCount++;
+  if (self.attachmentCleanupHandler) {
+    self.attachmentCleanupHandler(self, promise);
+    return;
+  }
+  [super performAttachmentCleanupForPromise:promise];
+}
+@end
+
 @interface FBSession (SystemApplicationTests)
 - (BOOL)fb_isApplicationSameAsSystemApp:(XCUIApplication *)application timeout:(NSTimeInterval)timeout;
 @end
@@ -265,4 +302,104 @@ - (void)testSystemApplicationLookupTimeoutSkipsTermination
   }
 }
 
+- (FBScreenRecordingPromise *)makeScreenRecordingPromiseWithUUID:(NSUUID *)uuid
+{
+  FBNativeScreenRecordingPromiseDouble *nativePromise = [FBNativeScreenRecordingPromiseDouble new];
+  nativePromise.identifier = uuid;
+  return [[FBScreenRecordingPromise alloc] initWithNativePromise:nativePromise];
+}
+
+- (void)testScreenRecordingContainerResetClearsStateWhenAttachmentCleanupThrows
+{
+  FBTestableScreenRecordingContainer *container = [FBTestableScreenRecordingContainer new];
+  NSUUID *uuid = [NSUUID UUID];
+  FBScreenRecordingPromise *promise = [self makeScreenRecordingPromiseWithUUID:uuid];
+  [container storeScreenRecordingPromise:promise fps:24 codec:1];
+  container.attachmentCleanupHandler = ^(FBScreenRecordingContainer *targetContainer, FBScreenRecordingPromise *cleanedPromise) {
+    XCTAssertNotNil(targetContainer);
+    XCTAssertEqual(cleanedPromise, promise);
+    @throw [NSException exceptionWithName:@"CleanupFailure" reason:@"Attachment failed" userInfo:nil];
+  };
+
+  XCTAssertNoThrow([container reset]);
+  XCTAssertEqual(container.cleanupCallCount, 1U);
+  XCTAssertNil(container.screenRecordingPromise);
+  XCTAssertEqual(container.fps, 0U);
+  XCTAssertEqual(container.codec, 0LL);
+  XCTAssertNil(container.startedAt);
+  XCTAssertNil([container toDictionary]);
+}
+
+- (void)testScreenRecordingContainerClearsStateBeforeAttachmentCleanupRuns
+{
+  FBTestableScreenRecordingContainer *container = [FBTestableScreenRecordingContainer new];
+  FBScreenRecordingPromise *promise = [self makeScreenRecordingPromiseWithUUID:[NSUUID UUID]];
+  [container storeScreenRecordingPromise:promise fps:24 codec:1];
+
+  __block FBScreenRecordingPromise *promiseDuringCleanup = promise;
+  __block NSDictionary *dictionaryDuringCleanup = @{};
+  __block NSUInteger fpsDuringCleanup = 24;
+  container.attachmentCleanupHandler = ^(FBScreenRecordingContainer *targetContainer, FBScreenRecordingPromise *cleanedPromise) {
+    XCTAssertEqual(cleanedPromise, promise);
+    promiseDuringCleanup = targetContainer.screenRecordingPromise;
+    dictionaryDuringCleanup = [targetContainer toDictionary];
+    fpsDuringCleanup = targetContainer.fps;
+  };
+
+  [container reset];
+  XCTAssertEqual(container.cleanupCallCount, 1U);
+  XCTAssertNil(promiseDuringCleanup);
+  XCTAssertNil(dictionaryDuringCleanup);
+  XCTAssertEqual(fpsDuringCleanup, 0U);
+}
+
+- (void)testScreenRecordingContainerPreservesConcurrentReplacementDuringCleanup
+{
+  FBTestableScreenRecordingContainer *container = [FBTestableScreenRecordingContainer new];
+  FBScreenRecordingPromise *firstPromise = [self makeScreenRecordingPromiseWithUUID:[NSUUID UUID]];
+  NSUUID *secondUUID = [NSUUID UUID];
+  FBScreenRecordingPromise *secondPromise = [self makeScreenRecordingPromiseWithUUID:secondUUID];
+  [container storeScreenRecordingPromise:firstPromise fps:24 codec:0];
+
+  container.attachmentCleanupHandler = ^(FBScreenRecordingContainer *targetContainer, FBScreenRecordingPromise *cleanedPromise) {
+    XCTAssertEqual(cleanedPromise, firstPromise);
+    [targetContainer storeScreenRecordingPromise:secondPromise fps:60 codec:1];
+  };
+
+  XCTAssertTrue([container resetIfPromiseIs:firstPromise]);
+  XCTAssertEqual(container.cleanupCallCount, 1U);
+  XCTAssertEqual(container.screenRecordingPromise, secondPromise);
+  XCTAssertEqual(container.fps, 60U);
+  XCTAssertEqual(container.codec, 1LL);
+  NSDictionary *metadata = [container toDictionary];
+  XCTAssertEqualObjects(metadata[@"fps"], @(60));
+  XCTAssertEqualObjects(metadata[@"codec"], @(1));
+  XCTAssertEqualObjects(metadata[@"uuid"], secondUUID.UUIDString);
+}
+
+- (void)testScreenRecordingContainerResetIfPromiseIsIgnoresStalePromise
+{
+  FBTestableScreenRecordingContainer *container = [FBTestableScreenRecordingContainer new];
+  NSUUID *activeUUID = [NSUUID UUID];
+  FBScreenRecordingPromise *activePromise = [self makeScreenRecordingPromiseWithUUID:activeUUID];
+  FBScreenRecordingPromise *stalePromise = [self makeScreenRecordingPromiseWithUUID:[NSUUID UUID]];
+  [container storeScreenRecordingPromise:activePromise fps:30 codec:1];
+
+  XCTAssertFalse([container resetIfPromiseIs:stalePromise]);
+  XCTAssertEqual(container.cleanupCallCount, 0U);
+  XCTAssertEqual(container.screenRecordingPromise, activePromise);
+  XCTAssertEqual(container.fps, 30U);
+  XCTAssertEqual(container.codec, 1LL);
+  XCTAssertEqualObjects([container toDictionary][@"uuid"], activeUUID.UUIDString);
+}
+
+- (void)testScreenRecordingContainerResetOnEmptyContainerDoesNotRunCleanup
+{
+  FBTestableScreenRecordingContainer *container = [FBTestableScreenRecordingContainer new];
+  [container reset];
+  XCTAssertEqual(container.cleanupCallCount, 0U);
+  XCTAssertNil(container.screenRecordingPromise);
+  XCTAssertNil([container toDictionary]);
+}
+
 @end