FirebaseMessaging: Mock FIRHeartbeatInfo in test (#7995)

Author: ncooke3

FirebaseCore/Sources/Private/FIRHeartbeatInfo.h

@@ -27,7 +27,7 @@ typedef NS_ENUM(NSInteger, FIRHeartbeatInfoCode) {
 };
 
 /**
- * Get heartbeat code requred for the sdk.
+ * Get heartbeat code required for the sdk.
  * @param heartbeatTag String representing the sdk heartbeat tag.
  * @return Heartbeat code indicating whether or not an sdk/global heartbeat
  * needs to be sent

FirebaseMessaging/Tests/UnitTests/FIRMessagingTokenOperationsTest.m

@@ -18,7 +18,6 @@
 
 #import "OCMock.h"
 
-#import <GoogleUtilities/GULHeartbeatDateStorage.h>
 #import "FirebaseCore/Sources/Private/FirebaseCoreInternal.h"
 #import "FirebaseInstallations/Source/Library/Private/FirebaseInstallationsInternal.h"
 #import "FirebaseMessaging/Sources/FIRMessagingConstants.h"
@@ -78,6 +77,7 @@ @interface FIRMessagingTokenOperationsTest : XCTestCase
 @property(strong, readonly, nonatomic) FIRMessagingCheckinService *checkinService;
 @property(strong, readonly, nonatomic) id mockCheckinService;
 @property(strong, readonly, nonatomic) id mockInstallations;
+@property(strong, readonly, nonatomic) id mockHeartbeatInfo;
 
 @property(strong, readonly, nonatomic) NSString *instanceID;
 
@@ -107,10 +107,9 @@ - (void)setUp {
   // Stub `FIRInstallations` to avoid using a real object.
   [self stubInstallations];
 
-  NSString *const kHeartbeatStorageFile = @"HEARTBEAT_INFO_STORAGE";
-  GULHeartbeatDateStorage *dataStorage =
-      [[GULHeartbeatDateStorage alloc] initWithFileName:kHeartbeatStorageFile];
-  [[NSFileManager defaultManager] removeItemAtURL:[dataStorage fileURL] error:nil];
+  // `FIRMessagingTokenFetchOperation` uses `FIRHeartbeatInfo` to retrieve a heartbeat code.
+  // Stub `FIRHeartbeatInfo` to avoid using a real object.
+  [self stubHeartbeatInfo];
 }
 
 - (void)tearDown {
@@ -120,6 +119,7 @@ - (void)tearDown {
   _checkinService = nil;
   _mockTokenStore = nil;
   [_mockInstallations stopMocking];
+  [_mockHeartbeatInfo stopMocking];
 }
 
 - (void)testThatTokenOperationsAuthHeaderStringMatchesCheckin {
@@ -390,7 +390,9 @@ - (void)testTokenFetchOperationFirebaseUserAgentAndHeartbeatHeader {
                     XCTAssertEqualObjects(userAgentValue, [FIRApp firebaseUserAgent]);
                     NSString *heartBeatCode =
                         sentRequest.allHTTPHeaderFields[kFIRMessagingFirebaseHeartbeatKey];
-                    XCTAssertEqualObjects(heartBeatCode, @"3");
+                    // It is expected that both the SDK and global heartbeat are requested.
+                    XCTAssertEqual(heartBeatCode.integerValue, FIRHeartbeatInfoCodeCombined,
+                                   @"Heartbeat storage info needed to be updated but was not.");
                     [completionExpectation fulfill];
 
                     return YES;
@@ -438,4 +440,10 @@ - (void)stubInstallations {
   OCMStub([_mockInstallations authTokenWithCompletion:authTokenWithCompletionArg]);
 }
 
+- (void)stubHeartbeatInfo {
+  _mockHeartbeatInfo = OCMClassMock([FIRHeartbeatInfo class]);
+  OCMStub([_mockHeartbeatInfo heartbeatCodeForTag:@"fire-iid"])
+      .andReturn(FIRHeartbeatInfoCodeCombined);
+}
+
 @end