Deflake SNTApplicationTest by tracking subscriptions to specific even…

…t types (#614) * Switch to waitForExpectations in tests * Change mock es_subscribe to note specific events we're ready for
google · Sep 27, 2021 · 81049db · 81049db
1 parent c110245
commit 81049db
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 61 deletions.
diff --git a/Source/santad/EventProviders/EndpointSecurityTestUtil.h b/Source/santad/EventProviders/EndpointSecurityTestUtil.h
@@ -44,7 +44,7 @@ typedef void (^ESCallback)(ESResponse *_Nonnull);
 
 // Singleton wrapper around all of the kernel-level EndpointSecurity framework functions.
 @interface MockEndpointSecurity : NSObject
-@property BOOL subscribed;
+@property NSMutableArray *_Nonnull subscriptions;
 - (void)reset;
 - (void)registerResponseCallback:(ESCallback _Nonnull)callback;
 - (void)triggerHandler:(es_message_t *_Nonnull)msg;

diff --git a/Source/santad/EventProviders/EndpointSecurityTestUtil.mm b/Source/santad/EventProviders/EndpointSecurityTestUtil.mm
@@ -100,17 +100,23 @@ - (instancetype)init {
   self = [super init];
   if (self) {
     _responseCallbacks = [NSMutableArray array];
-    _subscribed = YES;
+    _subscriptions = [NSMutableArray arrayWithCapacity:ES_EVENT_TYPE_LAST];
+    [self resetSubscriptions];
   }
   return self;
 };
 
+- (void)resetSubscriptions {
+  for (size_t i = 0; i < ES_EVENT_TYPE_LAST; i++) {
+    _subscriptions[i] = @NO;
+  }
+}
+
 - (void)reset {
   @synchronized(self) {
     [self.responseCallbacks removeAllObjects];
     self.handler = nil;
     self.client = nil;
-    self.subscribed = NO;
   }
 };
 
@@ -148,6 +154,16 @@ - (es_respond_result_t)respond_auth_result:(const es_message_t *_Nonnull)msg
   return ES_RESPOND_RESULT_SUCCESS;
 };
 
+- (void)setSubscriptions:(const es_event_type_t *_Nonnull)events
+             event_count:(uint32_t)event_count
+                   value:(NSNumber *)value {
+  @synchronized(self) {
+    for (size_t i = 0; i < event_count; i++) {
+      self.subscriptions[events[i]] = value;
+    }
+  }
+}
+
 + (instancetype _Nonnull)mockEndpointSecurity {
   static MockEndpointSecurity *sharedES;
   static dispatch_once_t onceToken;
@@ -194,14 +210,18 @@ es_respond_result_t es_respond_auth_result(es_client_t *_Nonnull client,
 API_UNAVAILABLE(ios, tvos, watchos)
 es_return_t es_subscribe(es_client_t *_Nonnull client, const es_event_type_t *_Nonnull events,
                          uint32_t event_count) {
-  [MockEndpointSecurity mockEndpointSecurity].subscribed = YES;
+  [[MockEndpointSecurity mockEndpointSecurity] setSubscriptions:events
+                                                    event_count:event_count
+                                                          value:@YES];
   return ES_RETURN_SUCCESS;
 }
 API_AVAILABLE(macos(10.15))
 API_UNAVAILABLE(ios, tvos, watchos)
 es_return_t es_unsubscribe(es_client_t *_Nonnull client, const es_event_type_t *_Nonnull events,
                            uint32_t event_count) {
-  [MockEndpointSecurity mockEndpointSecurity].subscribed = NO;
+  [[MockEndpointSecurity mockEndpointSecurity] setSubscriptions:events
+                                                    event_count:event_count
+                                                          value:@NO];
 
   return ES_RETURN_SUCCESS;
 };
diff --git a/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm b/Source/santad/EventProviders/SNTEndpointSecurityManagerTest.mm
@@ -71,14 +71,7 @@ - (void)testDenyOnTimeout {
 
   [mockES triggerHandler:m.message];
 
-  [self waitForExpectationsWithTimeout:30.0
-                               handler:^(NSError *error) {
-                                 if (error) {
-                                   XCTFail(@"Santa auth test timed out without receiving two "
-                                           @"events. Instead, had error: %@",
-                                           error);
-                                 }
-                               }];
+  [self waitForExpectations:@[ expectation ] timeout:60.0];
 
   for (ESResponse *resp in events) {
     XCTAssertEqual(
@@ -116,12 +109,7 @@ - (void)testDeleteRulesDB {
 
     [mockES triggerHandler:m.message];
 
-    [self waitForExpectationsWithTimeout:30.0
-                                 handler:^(NSError *error) {
-                                   if (error) {
-                                     XCTFail(@"Santa auth test timed out with error: %@", error);
-                                   }
-                                 }];
+    [self waitForExpectations:@[ expectation ] timeout:60.0];
 
     XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue,
                    @"Incorrect handling of delete of %@", testPath);
@@ -152,12 +140,8 @@ - (void)testSkipOtherESEvents {
   }];
 
   [mockES triggerHandler:m.message];
-  [self waitForExpectationsWithTimeout:30.0
-                               handler:^(NSError *error) {
-                                 if (error) {
-                                   XCTFail(@"Santa auth test timed out with error: %@", error);
-                                 }
-                               }];
+
+  [self waitForExpectations:@[ expectation ] timeout:60.0];
 
   XCTAssertEqual(got.result, ES_AUTH_RESULT_ALLOW);
 }
@@ -199,12 +183,7 @@ - (void)testRenameOverwriteRulesDB {
 
     [mockES triggerHandler:m.message];
 
-    [self waitForExpectationsWithTimeout:30.0
-                                 handler:^(NSError *error) {
-                                   if (error) {
-                                     XCTFail(@"Santa auth test timed out with error: %@", error);
-                                   }
-                                 }];
+    [self waitForExpectations:@[ expectation ] timeout:60.0];
 
     XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue,
                    @"Incorrect handling of rename of %@", testPath);
@@ -254,12 +233,8 @@ - (void)testRenameRulesDB {
 
     [mockES triggerHandler:m.message];
 
-    [self waitForExpectationsWithTimeout:30.0
-                                 handler:^(NSError *error) {
-                                   if (error) {
-                                     XCTFail(@"Santa auth test timed out with error: %@", error);
-                                   }
-                                 }];
+    [self waitForExpectations:@[ expectation ] timeout:60.0];
+
     XCTAssertEqual(got.result, [testCases objectForKey:testPath].intValue,
                    @"Incorrect handling of rename of %@", testPath);
 

diff --git a/Source/santad/SNTApplicationTest.m b/Source/santad/SNTApplicationTest.m
@@ -52,34 +52,26 @@ - (void)checkBinaryExecution:(NSString *)binaryName
   SNTApplication *app = [[SNTApplication alloc] init];
   [app start];
 
-  // es events will start flowing in as soon as es_subscribe is called, regardless
-  // of whether we're ready or not for it.
   XCTestExpectation *santaInit =
     [self expectationWithDescription:@"Wait for Santa to subscribe to EndpointSecurity"];
 
   dispatch_async(dispatch_get_global_queue(QOS_CLASS_BACKGROUND, 0), ^{
-    while (!mockES.subscribed)
+    while ([mockES.subscriptions[ES_EVENT_TYPE_AUTH_EXEC] isEqualTo:@NO])
       ;
     [santaInit fulfill];
   });
 
-  [self waitForExpectationsWithTimeout:30.0
-                               handler:^(NSError *error) {
-                                 if (error) {
-                                   XCTFail(@"Santa's subscription to EndpointSecurity timed out "
-                                           @"with error: %@",
-                                           error);
-                                 }
-                               }];
+  // Ugly hack to deflake the test and allow listenForDecisionRequests to install the correct
+  // decision callback.
+  sleep(1);
+  [self waitForExpectations:@[ santaInit ] timeout:10.0];
 
   XCTestExpectation *expectation =
     [self expectationWithDescription:@"Wait for santa's Auth dispatch queue"];
   __block ESResponse *got = nil;
   [mockES registerResponseCallback:^(ESResponse *r) {
-    @synchronized(self) {
-      got = r;
-      [expectation fulfill];
-    }
+    got = r;
+    [expectation fulfill];
   }];
 
   NSString *binaryPath = [NSString pathWithComponents:@[ testPath, binaryName ]];
@@ -95,15 +87,7 @@ - (void)checkBinaryExecution:(NSString *)binaryName
 
   [mockES triggerHandler:msg.message];
 
-  [self
-    waitForExpectationsWithTimeout:30.0
-                           handler:^(NSError *error) {
-                             if (error) {
-                               XCTFail(
-                                 @"Santa auth test on binary \"%@/%@\" timed out with error: %@",
-                                 testPath, binaryName, error);
-                             }
-                           }];
+  [self waitForExpectations:@[ expectation ] timeout:10.0];
 
   XCTAssertEqual(got.result, wantResult, @"received unexpected ES response on executing \"%@/%@\"",
                  testPath, binaryName);