Tests: Fix MetricServiceTest compatible with public OCMock

google · Nov 8, 2021 · 0a314cb · 0a314cb
1 parent ad6e03e
commit 0a314cb
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 40 deletions.
diff --git a/Source/santametricservice/SNTMetricServiceTest.m b/Source/santametricservice/SNTMetricServiceTest.m
@@ -43,15 +43,6 @@ - (void)setUp {
 
 - (void)tearDown {
   [self.mockConfigurator stopMocking];
-  if (self.mockSessionDataTask != nil) {
-    [self.mockSessionDataTask stopMocking];
-  }
-  if (self.mockSession != nil) {
-    [self.mockSession stopMocking];
-  }
-  if (self.mockMOLAuthenticatingURLSession != nil) {
-    [self.mockMOLAuthenticatingURLSession stopMocking];
-  }
 
   // delete the temp dir
   [[NSFileManager defaultManager] removeItemAtPath:self.tempDir error:NULL];
@@ -146,28 +137,25 @@ - (void)testWritingJSONOverHTTP {
                                HTTPVersion:@"HTTP/1.1"
                               headerFields:@{@"content-type" : @"application/json"}];
 
-  typedef void (^dataTaskCompletion)(NSData *, NSURLResponse *, NSError *);
-  typedef id (^completionHandler)(NSURLSession *, NSURLRequest *, dataTaskCompletion);
-
-  __block dataTaskCompletion passedBlock;
+  __unsafe_unretained __block void (^passedBlock)(NSData *, NSURLResponse *, NSError *);
 
   XCTestExpectation *responseCallback =
     [[XCTestExpectation alloc] initWithDescription:@"ensure writer passed JSON"];
 
   // stub out session to call completion handler immediately.
-  OCMStub([(NSURLSessionDataTask *)self.mockSessionDataTask resume])
-    .andDo(^void(NSURLSessionDataTask *unused) {
+  OCMStub([(NSURLSessionDataTask *)self.mockSessionDataTask resume]).andDo(^(NSInvocation *inv) {
+    if (passedBlock) {
       passedBlock(nil, response, nil);
-      [responseCallback fulfill];
-    });
+    }
+    [responseCallback fulfill];
+  });
 
   // stub out NSURLSession to assign our completion handler and return our mock
   OCMStub([self.mockSession dataTaskWithRequest:[OCMArg any] completionHandler:[OCMArg any]])
-    .andDo(
-      ^id(NSURLSession *localSelf, NSURLRequest *request, dataTaskCompletion completionHandler) {
-        passedBlock = completionHandler;
-        return nil;
-      })
+    .andDo(^(NSInvocation *inv) {
+      [inv retainArguments];
+      [inv getArgument:&passedBlock atIndex:3];
+    })
     .andReturn(self.mockSessionDataTask);
 
   SNTMetricService *service = [[SNTMetricService alloc] init];

diff --git a/Source/santametricservice/Writers/SNTMetricHTTPWriterTest.m b/Source/santametricservice/Writers/SNTMetricHTTPWriterTest.m
@@ -28,26 +28,27 @@ - (void)setUp {
   self.httpWriter = [[SNTMetricHTTPWriter alloc] init];
   self.mockResponses = [[NSMutableArray alloc] init];
 
-  typedef void (^dataTaskCompletion)(NSData *, NSURLResponse *, NSError *);
-  __block dataTaskCompletion completionHandler;
-
-  OCMStub([(NSURLSessionDataTask *)self.mockSessionDataTask resume])
-    .andDo(^void(NSURLSessionDataTask *unused) {
-      NSDictionary *responseValue = self.mockResponses[0];
-      if (responseValue != nil) {
-        completionHandler(responseValue[@"data"], responseValue[@"response"],
-                          responseValue[@"error"]);
-        [self.mockResponses removeObjectAtIndex:0];
-      } else {
-        XCTFail(@"mockResponses set to zero");
-      }
-    });
+  __block void (^completionHandler)(NSData *, NSURLResponse *, NSError *);
+
+  void (^getCompletionHandler)(NSInvocation *) = ^(NSInvocation *invocation) {
+    [invocation getArgument:&completionHandler atIndex:3];
+  };
+
+  void (^callCompletionHandler)(NSInvocation *) = ^(NSInvocation *invocation) {
+    NSDictionary *responseValue = self.mockResponses[0];
+    if (responseValue != nil) {
+      completionHandler(responseValue[@"data"], responseValue[@"response"],
+                        responseValue[@"error"]);
+      [self.mockResponses removeObjectAtIndex:0];
+    } else {
+      XCTFail(@"mockResponses set to zero");
+    }
+  };
+
+  OCMStub([(NSURLSessionDataTask *)self.mockSessionDataTask resume]).andDo(callCompletionHandler);
 
   OCMStub([self.mockSession dataTaskWithRequest:[OCMArg any] completionHandler:[OCMArg any]])
-    .andDo(^id(NSURLSession *localSelf, NSURLRequest *request, dataTaskCompletion completion) {
-      completionHandler = completion;
-      return nil;
-    })
+    .andDo(getCompletionHandler)
     .andReturn(self.mockSessionDataTask);
 }