IID tests improvements (#3067)

* Use bridged __weak reference instead CFRetain/CFRelease
* Use SecKeyRef instead of FIRInstanceIDKeyPair for the memory management test
* Use SecKeyGeneratePair instead of SecKeyCreateRandomKey available since iOS 10 only
* FIRInstanceIDKeyPairStoreTest wait until all keychain operations finished before finishing a test
* FIRInstanceIDTokenManagerTest: cleanup in tearDown; fix expectation to produce an error nor a crash.
* Use properties instead of ivars FIRInstanceIDTokenManagerTest
* FIRInstanceIDTokenManager fix OCMPartialMock() usage.

Author: maksymmalyhin

Example/InstanceID/Tests/FIRInstanceIDKeyPairMigrationTest.m

@@ -70,6 +70,10 @@ - (void)tearDown {
   [super tearDown];
   NSError *error = nil;
   [self.keyPairStore removeKeyPairCreationTimePlistWithError:&error];
+
+  // TODO: Real FIRInstanceIDKeychain should not be used for the tests, it should be mocked instead.
+  // Drain Keychain private queue before exiting.
+  [[FIRInstanceIDKeychain sharedInstance] itemWithQuery:@{}];
 }
 
 - (void)testMigrationDataIfLegacyKeyPairsNotExist {
@@ -139,33 +143,49 @@ - (void)testMigrationIfLegacyKeyPairsExist {
   [self waitForExpectationsWithTimeout:1 handler:nil];
 }
 
-// Disabling test for now. We need to find a flake free way to insure the publicKeyRef is retained.
-#ifdef DISABLED
 - (void)testUpdateKeyRefWithTagRetainsAndReleasesKeyRef {
-  SecKeyRef publicKeyRef;
+  __weak id weakKeyRef;
 
+  // Use a local autorelease pool to make sure any autorelease objects allocated will be released.
   @autoreleasepool {
-    NSString *legacyPublicKeyTag =
-        FIRInstanceIDLegacyPublicTagWithSubtype(kFIRInstanceIDKeyPairSubType);
-    NSString *legacyPrivateKeyTag =
-        FIRInstanceIDLegacyPrivateTagWithSubtype(kFIRInstanceIDKeyPairSubType);
-    FIRInstanceIDKeyPair *keyPair =
-        [[FIRInstanceIDKeychain sharedInstance] generateKeyPairWithPrivateTag:legacyPrivateKeyTag
-                                                                    publicTag:legacyPublicKeyTag];
-    XCTAssertTrue([keyPair isValid]);
-
-    publicKeyRef = keyPair.publicKey;
-
+    SecKeyRef keyRef = [self generateKeyRef];
+    weakKeyRef = (__bridge id)(keyRef);
     XCTestExpectation *completionExpectation =
         [self expectationWithDescription:@"completionExpectation"];
-    [self.keyPairStore updateKeyRef:keyPair.publicKey
+    [self.keyPairStore updateKeyRef:keyRef
                             withTag:@"test"
                             handler:^(NSError *error) {
                               [completionExpectation fulfill];
                             }];
+
+    // Release locally allocated CoreFoundation object.
+    CFRelease(keyRef);
   }
+
+  // Should be still alive until execution finished
+  XCTAssertNotNil(weakKeyRef);
   [self waitForExpectationsWithTimeout:0.5 handler:NULL];
+
+  // Should be released once finished
+  // The check below is flaky for build under DEBUG (petentially due to ARC specifics).
+  // Comment it so far as not-so-important one.
+  //  XCTAssertNil(weakKeyRef);
+}
+
+- (SecKeyRef)generateKeyRef {
+  NSDictionary *keyAttributes = @{(__bridge id)kSecAttrIsPermanent : @YES};
+  NSDictionary *keyPairAttributes = @{
+    (__bridge id)kSecAttrKeyType : (__bridge id)kSecAttrKeyTypeRSA,
+    (__bridge id)kSecAttrLabel : @"[FIRInstanceIDKeyPairMigrationTest generateKeyRef]",
+    (__bridge id)kSecAttrKeySizeInBits : @(2048),
+    (__bridge id)kSecPrivateKeyAttrs : keyAttributes,
+    (__bridge id)kSecPublicKeyAttrs : keyAttributes,
+  };
+
+  SecKeyRef publicKey = NULL;
+  SecKeyGeneratePair((__bridge CFDictionaryRef)keyPairAttributes, &publicKey, NULL);
+
+  return publicKey;
 }
-#endif
 
 @end

Example/InstanceID/Tests/FIRInstanceIDKeyPairStoreTest.m

@@ -62,10 +62,17 @@ - (void)setUp {
 }
 
 - (void)tearDown {
-  [super tearDown];
   NSError *error = nil;
   [self.keyPairStore removeKeyPairCreationTimePlistWithError:&error];
-  [self.keyPairStore deleteSavedKeyPairWithSubtype:kFIRInstanceIDKeyPairSubType handler:nil];
+
+  XCTestExpectation *queueDrained = [self expectationWithDescription:@"drainKeychainQueue"];
+  [self.keyPairStore deleteSavedKeyPairWithSubtype:kFIRInstanceIDKeyPairSubType
+                                           handler:^(NSError *error) {
+                                             [queueDrained fulfill];
+                                           }];
+  [self waitForExpectations:@[ queueDrained ] timeout:10];
+
+  [super tearDown];
 }
 
 /**
@@ -125,23 +132,23 @@ - (void)testResetIdentity {
   XCTAssertNil(error);
   NSString *iid1 = FIRInstanceIDAppIdentity(keyPair);
 
-  [self.keyPairStore
-      deleteSavedKeyPairWithSubtype:kFIRInstanceIDKeyPairSubType
-                            handler:^(NSError *error) {
-                              XCTAssertNil(error);
-                              [self.keyPairStore removeKeyPairCreationTimePlistWithError:&error];
-                              XCTAssertNil(error);
-
-                              // regenerate instance-id
-                              FIRInstanceIDKeyPair *keyPair =
-                                  [self.keyPairStore loadKeyPairWithError:&error];
-                              XCTAssertNil(error);
-                              NSString *iid2 = FIRInstanceIDAppIdentity(keyPair);
+  [self.keyPairStore deleteSavedKeyPairWithSubtype:kFIRInstanceIDKeyPairSubType
+                                           handler:^(NSError *error) {
+                                             XCTAssertNil(error);
+                                             [identityResetExpectation fulfill];
+                                           }];
 
-                              XCTAssertNotEqualObjects(iid1, iid2);
-                              [identityResetExpectation fulfill];
-                            }];
   [self waitForExpectationsWithTimeout:5 handler:nil];
+
+  [self.keyPairStore removeKeyPairCreationTimePlistWithError:&error];
+  XCTAssertNil(error);
+
+  // regenerate instance-id
+  FIRInstanceIDKeyPair *keyPair2 = [self.keyPairStore loadKeyPairWithError:&error];
+  XCTAssertNil(error);
+  NSString *iid2 = FIRInstanceIDAppIdentity(keyPair2);
+
+  XCTAssertNotEqualObjects(iid1, iid2);
 }
 
 /**
@@ -236,4 +243,5 @@ - (void)testDeleteKeyPair {
                             }];
   [self waitForExpectationsWithTimeout:1 handler:nil];
 }
+
 @end

Example/InstanceID/Tests/FIRInstanceIDKeyPairTest.m

@@ -65,8 +65,13 @@ - (void)testValidKeychain {
   XCTAssertNotNil(keypair.privateKeyData);
   XCTAssertNotNil(FIRInstanceIDAppIdentity(keypair));
 
+  XCTestExpectation *keyPairDeleted = [self expectationWithDescription:@"keyPairDeleted"];
   [FIRInstanceIDKeyPairStore deleteKeyPairWithPrivateTag:kKeyPairPrivateTag
                                                publicTag:kKeyPairPublicTag
-                                                 handler:nil];
+                                                 handler:^(NSError *error) {
+                                                   [keyPairDeleted fulfill];
+                                                 }];
+
+  [self waitForExpectations:@[ keyPairDeleted ] timeout:1.0];
 }
 @end

Example/InstanceID/Tests/FIRInstanceIDResultTest.m

@@ -45,13 +45,12 @@ @implementation FIRInstanceIDResultTest
 
 - (void)setUp {
   [super setUp];
-  _instanceID = [[FIRInstanceID alloc] initPrivately];
-  [_instanceID start];
-  _mockInstanceID = OCMPartialMock(_instanceID);
+  _mockInstanceID = OCMClassMock([FIRInstanceID class]);
 }
 
 - (void)tearDown {
   [_mockInstanceID stopMocking];
+  _mockInstanceID = nil;
   [super tearDown];
 }