[Darwin] Fix flaky subscription pool test (project-chip#35606)

Sarthak-Shaha · Sep 17, 2024 · cc5ea19 · cc5ea19
1 parent c06a0be
commit cc5ea19
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 17 deletions.
diff --git a/src/darwin/Framework/CHIP/MTRDevice_Concrete.mm b/src/darwin/Framework/CHIP/MTRDevice_Concrete.mm
@@ -1167,8 +1167,7 @@ - (void)_scheduleSubscriptionPoolWork:(dispatch_block_t)workBlock inNanoseconds:
         return;
     }
 
-    // Wait the required amount of time, then put it in the subscription pool to wait additionally for a spot, if needed
-    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, inNanoseconds), self.queue, ^{
+    dispatch_block_t workBlockToQueue = ^{
         // In the case where a resubscription triggering event happened and already established, running the work block should result in a no-op
         MTRAsyncWorkItem * workItem = [[MTRAsyncWorkItem alloc] initWithQueue:self.queue];
         [workItem setReadyHandler:^(id _Nonnull context, NSInteger retryCount, MTRAsyncWorkCompletionBlock _Nonnull completion) {
@@ -1199,7 +1198,15 @@ - (void)_scheduleSubscriptionPoolWork:(dispatch_block_t)workBlock inNanoseconds:
         }];
         [self->_deviceController.concurrentSubscriptionPool enqueueWorkItem:workItem description:description];
         MTR_LOG("%@ - enqueued in the subscription pool", self);
-    });
+    };
+
+    if (inNanoseconds > 0) {
+        // Wait the required amount of time, then put it in the subscription pool to wait additionally for a spot
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, inNanoseconds), self.queue, workBlockToQueue);
+    } else {
+        // Put in subscription pool directly if there is no wait time
+        workBlockToQueue();
+    }
 }
 
 - (void)_handleResubscriptionNeededWithDelay:(NSNumber *)resubscriptionDelayMs

diff --git a/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m b/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m
@@ -2597,20 +2597,17 @@ - (void)doTestSubscriptionPoolWithSize:(NSInteger)subscriptionPoolSize deviceOnb
 
     // Create the base device to attempt to read from the 5th device
     __auto_type * baseDeviceReadExpectation = [self expectationWithDescription:@"BaseDevice read"];
-    // Dispatch async to get around XCTest, so that this runs after the above devices queue their subscriptions
-    dispatch_async(queue, ^{
-        __auto_type * baseDevice = [MTRBaseDevice deviceWithNodeID:@(105) controller:controller];
-        __auto_type * onOffCluster = [[MTRBaseClusterOnOff alloc] initWithDevice:baseDevice endpointID:@(1) queue:queue];
-        [onOffCluster readAttributeOnOffWithCompletion:^(NSNumber * value, NSError * _Nullable error) {
-            XCTAssertNil(error);
-            // We expect the device to be off.
-            XCTAssertEqualObjects(value, @(0));
-            [baseDeviceReadExpectation fulfill];
-            os_unfair_lock_lock(&counterLock);
-            baseDeviceReadCompleted = YES;
-            os_unfair_lock_unlock(&counterLock);
-        }];
-    });
+    __auto_type * baseDevice = [MTRBaseDevice deviceWithNodeID:@(105) controller:controller];
+    __auto_type * onOffCluster = [[MTRBaseClusterOnOff alloc] initWithDevice:baseDevice endpointID:@(1) queue:queue];
+    [onOffCluster readAttributeOnOffWithCompletion:^(NSNumber * value, NSError * _Nullable error) {
+        XCTAssertNil(error);
+        // We expect the device to be off.
+        XCTAssertEqualObjects(value, @(0));
+        [baseDeviceReadExpectation fulfill];
+        os_unfair_lock_lock(&counterLock);
+        baseDeviceReadCompleted = YES;
+        os_unfair_lock_unlock(&counterLock);
+    }];
 
     // Make the wait time depend on pool size and device count (can expand number of devices in the future)
     NSArray * expectationsToWait = [subscriptionExpectations.allValues arrayByAddingObject:baseDeviceReadExpectation];