1 files changed, 11 insertions, 4 deletions
diff --git a/Example/Messaging/Tests/FIRMessagingTest.m b/Example/Messaging/Tests/FIRMessagingTest.m
index 61ff136..92cf0f5 100644
--- a/Example/Messaging/Tests/FIRMessagingTest.m
+++ b/Example/Messaging/Tests/FIRMessagingTest.m
@@ -17,6 +17,8 @@
 #import <XCTest/XCTest.h>
 
 #import <OCMock/OCMock.h>
+
+#import <FirebaseCore/FIRAppInternal.h>
 #import <FirebaseInstanceID/FirebaseInstanceID.h>
 
 #import "FIRMessaging.h"
@@ -44,6 +46,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 @property(nonatomic, readonly, strong) FIRMessaging *messaging;
 @property(nonatomic, readwrite, strong) id mockMessaging;
 @property(nonatomic, readwrite, strong) id mockInstanceID;
+@property(nonatomic, readwrite, strong) id mockFirebaseApp;
 
 @end
 
@@ -51,6 +54,9 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 
 - (void)setUp {
   [super setUp];
+  _mockFirebaseApp = OCMClassMock([FIRApp class]);
+  OCMStub([_mockFirebaseApp defaultApp]).andReturn(_mockFirebaseApp);
+
   _messaging = [[FIRMessaging alloc] initWithInstanceID:[FIRInstanceID instanceID]
                                            userDefaults:[NSUserDefaults standardUserDefaults]];
   _mockMessaging = OCMPartialMock(self.messaging);
@@ -63,6 +69,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 - (void)tearDown {
   [_mockMessaging stopMocking];
   [_mockInstanceID stopMocking];
+  [_mockFirebaseApp stopMocking];
   [super tearDown];
 }
 
@@ -76,7 +83,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 }
 
 - (void)testAutoInitEnableFlagOverrideGlobalTrue {
-  OCMStub([self.mockMessaging isGlobalAutomaticDataCollectionEnabled]).andReturn(YES);
+  OCMStub([_mockFirebaseApp isAutomaticDataCollectionEnabled]).andReturn(YES);
   id bundleMock = OCMPartialMock([NSBundle mainBundle]);
   OCMStub([bundleMock objectForInfoDictionaryKey:kFIRMessagingPlistAutoInitEnabled]).andReturn(nil);
   XCTAssertTrue(self.messaging.isAutoInitEnabled);
@@ -87,7 +94,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 }
 
 - (void)testAutoInitEnableFlagOverrideGlobalFalse {
-  OCMStub([self.mockMessaging isGlobalAutomaticDataCollectionEnabled]).andReturn(YES);
+  OCMStub([_mockFirebaseApp isAutomaticDataCollectionEnabled]).andReturn(YES);
   id bundleMock = OCMPartialMock([NSBundle mainBundle]);
   OCMStub([bundleMock objectForInfoDictionaryKey:kFIRMessagingPlistAutoInitEnabled]).andReturn(nil);
   XCTAssertTrue(self.messaging.isAutoInitEnabled);
@@ -98,7 +105,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 }
 
 - (void)testAutoInitEnableGlobalDefaultTrue {
-  OCMStub([self.mockMessaging isGlobalAutomaticDataCollectionEnabled]).andReturn(YES);
+  OCMStub([_mockFirebaseApp isAutomaticDataCollectionEnabled]).andReturn(YES);
   id bundleMock = OCMPartialMock([NSBundle mainBundle]);
   OCMStub([bundleMock objectForInfoDictionaryKey:kFIRMessagingPlistAutoInitEnabled]).andReturn(nil);
 
@@ -107,7 +114,7 @@ extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
 }
 
 - (void)testAutoInitEnableGlobalDefaultFalse {
-  OCMStub([self.mockMessaging isGlobalAutomaticDataCollectionEnabled]).andReturn(NO);
+  OCMStub([_mockFirebaseApp isAutomaticDataCollectionEnabled]).andReturn(NO);
   id bundleMock = OCMPartialMock([NSBundle mainBundle]);
   OCMStub([bundleMock objectForInfoDictionaryKey:kFIRMessagingPlistAutoInitEnabled]).andReturn(nil);