Add a unit test to verify behavior w/ URLSession:dataTask:didReceiveResponse:completionHandler:

.ruby-version

@@ -1 +1 @@
-ruby-2.7
+2.7.4

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTest.m

@@ -62,6 +62,35 @@
 
 @end
 
+@interface FPRNSURLSessionDelegateProxy : NSProxy <NSURLSessionDelegate> {
+  // The wrapped delegate object.
+  id _delegate;
+}
+
+/** @return an instance of the delegate proxy. */
+- (instancetype)initWithDelegate:(id <NSURLSessionDelegate>)delegate;
+
+@end
+
+@implementation FPRNSURLSessionDelegateProxy
+
+- (instancetype)initWithDelegate:(id)delegate {
+  if (self) {
+    _delegate = delegate;
+  }
+  return self;
+}
+
+- (NSMethodSignature *)methodSignatureForSelector:(SEL)selector {
+  return [_delegate methodSignatureForSelector:selector];
+}
+
+- (void)forwardInvocation:(NSInvocation *)invocation {
+  [invocation invokeWithTarget:_delegate];
+}
+
+@end
+
 @interface FPRNSURLSessionInstrumentTest : FPRTestCase
 
 /** Test server to create connections to. */
@@ -403,7 +432,7 @@
         [networkTrace
             timeIntervalBetweenCheckpointState:FPRNetworkTraceCheckpointStateInitiated
                                       andState:FPRNetworkTraceCheckpointStateRequestCompleted] > 0);
     XCTAssertNil([FPRNetworkTrace networkTraceFromObject:uploadTask]);
   }];
   [instrument deregisterInstrumentors];
 }
@@ -452,8 +481,8 @@
     XCTAssertNotNil([FPRNetworkTrace networkTraceFromObject:dataTask]);
     [self waitAndRunBlockAfterResponse:^(id self, GCDWebServerRequest *_Nonnull request,
                                          GCDWebServerResponse *_Nonnull response) {
       XCTAssertTrue(delegate.URLSessionDataTaskDidReceiveDataCalled);
       XCTAssertNil([FPRNetworkTrace networkTraceFromObject:dataTask]);
     }];
   }
   [instrument deregisterInstrumentors];
@@ -482,6 +511,60 @@
   [instrument deregisterInstrumentors];
 }
 
+/** Tests that the called delegate selector is wrapped and calls through. */
+- (void)testDelegateURLSessionDownloadDidReceiveResponseCompletionHandler {
+  FPRNSURLSessionInstrument *instrument;
+  NSURLSessionDataTask *dataTask;
+  @autoreleasepool {
+    instrument = [[FPRNSURLSessionInstrument alloc] init];
+    [instrument registerInstrumentors];
+    FPRNSURLSessionCompleteTestDelegate *delegate =
+        [[FPRNSURLSessionCompleteTestDelegate alloc] init];
+    NSURLSessionConfiguration *configuration =
+        [NSURLSessionConfiguration defaultSessionConfiguration];
+    NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                          delegate:delegate
+                                                     delegateQueue:nil];
+    NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testBigDownload"];
+    dataTask = [session dataTaskWithURL:URL];
+    [dataTask resume];
+    XCTAssertNotNil([FPRNetworkTrace networkTraceFromObject:dataTask]);
+    [self waitAndRunBlockAfterResponse:^(id self, GCDWebServerRequest *_Nonnull request,
+                                         GCDWebServerResponse *_Nonnull response) {
+      XCTAssertTrue(delegate.URLSessionDataTaskDidReceiveResponseCompletionHandlerCalled);
+      XCTAssertNil([FPRNetworkTrace networkTraceFromObject:dataTask]);
+    }];
+  }
+  [instrument deregisterInstrumentors];
+}
+
+/** Tests that the called delegate selector is wrapped and calls through. */
+- (void)testProxyDelegateURLSessionDownloadDidReceiveResponseCompletionHandler {
+  FPRNSURLSessionInstrument *instrument;
+  NSURLSessionDataTask *dataTask;
+  @autoreleasepool {
+    instrument = [[FPRNSURLSessionInstrument alloc] init];
+    [instrument registerInstrumentors];
+    FPRNSURLSessionCompleteTestDelegate *delegate =
+        [[FPRNSURLSessionCompleteTestDelegate alloc] init];
+    FPRNSURLSessionDelegateProxy *proxyDelegate = [[FPRNSURLSessionDelegateProxy alloc] initWithDelegate:delegate];
+    NSURLSessionConfiguration *configuration =
+        [NSURLSessionConfiguration defaultSessionConfiguration];
+    NSURLSession *session = [NSURLSession sessionWithConfiguration:configuration
+                                                          delegate:proxyDelegate
+                                                     delegateQueue:nil];
+    NSURL *URL = [self.testServer.serverURL URLByAppendingPathComponent:@"testBigDownload"];
+    dataTask = [session dataTaskWithURL:URL];
+    [dataTask resume];
+    XCTAssertNotNil([FPRNetworkTrace networkTraceFromObject:dataTask]);
+    [self waitAndRunBlockAfterResponse:^(id self, GCDWebServerRequest *_Nonnull request,
+                                         GCDWebServerResponse *_Nonnull response) {
+      XCTAssertTrue(delegate.URLSessionDataTaskDidReceiveResponseCompletionHandlerCalled);
+    }];
+  }
+  [instrument deregisterInstrumentors];
+}
+
 /** Tests that the called delegate selector is wrapped and calls through. */
 - (void)testDelegateURLSessionDownloadTaskDidResumeAtOffsetExpectedTotalBytes {
   FPRNSURLSessionInstrument *instrument = [[FPRNSURLSessionInstrument alloc] init];

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTestDelegates.h

@@ -52,6 +52,12 @@ NS_ASSUME_NONNULL_BEGIN
 /** Set to YES when URLSession:dataTask:didReceiveData: is called, used for testing. */
 @property(nonatomic) BOOL URLSessionDataTaskDidReceiveDataCalled;
 
+/** Set to YES when
+ * URLSession:dataTask:didReceiveResponse:completionHandler: is called, used
+ * for testing.
+ */
+@property(nonatomic) BOOL URLSessionDataTaskDidReceiveResponseCompletionHandlerCalled;
+
 /** Set to YES when URLSession:downloadTask:didFinishDownloadingToURL: is called, used for testing.
  */
 @property(nonatomic) BOOL URLSessionDownloadTaskDidFinishDownloadingToURLCalled;

FirebasePerformance/Tests/Unit/Instruments/FPRNSURLSessionInstrumentTestDelegates.m

@@ -95,6 +95,13 @@ - (void)URLSession:(NSURLSession *)session
   self.URLSessionDownloadTaskDidWriteDataTotalBytesWrittenTotalBytesCalled = YES;
 }
 
+- (void)URLSession:(NSURLSession *)session
+              dataTask:(NSURLSessionDataTask *)dataTask
+    didReceiveResponse:(NSURLResponse *)response
+     completionHandler:(void (^)(NSURLSessionResponseDisposition disposition))completionHandler {
+  self.URLSessionDataTaskDidReceiveResponseCompletionHandlerCalled = YES;
+}
+
 @end
 
 @interface FPRNSURLSessionTestDownloadDelegate ()