Initial

22 files changed, 4429 insertions, 0 deletions

diff --git a/Example/Messaging/Tests/FIRMessagingClientTest.m b/Example/Messaging/Tests/FIRMessagingClientTest.m
new file mode 100644
index 0000000..f36ec53
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingClientTest.m
@@ -0,0 +1,308 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+
+#import "FIRMessagingCheckinService.h"
+#import "FIRMessagingClient.h"
+#import "FIRMessagingConnection.h"
+#import "FIRMessagingDataMessageManager.h"
+#import "FIRMessagingFakeConnection.h"
+#import "FIRMessagingRegistrar.h"
+#import "FIRMessagingRmqManager.h"
+#import "FIRMessagingSecureSocket.h"
+#import "FIRMessagingUtilities.h"
+#import "NSError+FIRMessaging.h"
+
+#import "FIRReachabilityChecker.h"
+
+static NSString *const kFIRMessagingUserDefaultsSuite = @"FIRMessagingClientTestUserDefaultsSuite";
+
+static NSString *const kDeviceAuthId = @"123456";
+static NSString *const kSecretToken = @"56789";
+static NSString *const kDigest = @"com.google.digest";
+static NSString *const kVersionInfo = @"1.0";
+static NSString *const kSubscriptionID = @"abcdef-subscription-id";
+static NSString *const kDeletedSubscriptionID = @"deleted-abcdef-subscription-id";
+static NSString *const kFIRMessagingAppIDToken = @"1234xyzdef56789";
+static NSString *const kTopicToSubscribeTo = @"/topics/abcdef/hello-world";
+
+@interface FIRMessagingRegistrar ()
+
+@property(nonatomic, readwrite, strong) FIRMessagingCheckinService *checkinService;
+
+@end
+
+@interface FIRMessagingClient () <FIRMessagingConnectionDelegate>
+
+@property(nonatomic, readwrite, strong) FIRMessagingConnection *connection;
+@property(nonatomic, readwrite, strong) FIRMessagingRegistrar *registrar;
+
+@property(nonatomic, readwrite, assign) int64_t lastConnectedTimestamp;
+@property(nonatomic, readwrite, assign) int64_t lastDisconnectedTimestamp;
+@property(nonatomic, readwrite, assign) NSUInteger subscribeRetryCount;
+@property(nonatomic, readwrite, assign) NSUInteger connectRetryCount;
+
+- (NSTimeInterval)connectionTimeoutInterval;
+- (void)setupConnection;
+
+@end
+
+@interface FIRMessagingConnection () <FIRMessagingSecureSocketDelegate>
+
+@property(nonatomic, readwrite, strong) FIRMessagingSecureSocket *socket;
+
+- (void)setupConnectionSocket;
+- (void)connectToSocket:(FIRMessagingSecureSocket *)socket;
+- (NSTimeInterval)connectionTimeoutInterval;
+- (void)sendHeartbeatPing;
+
+@end
+
+@interface FIRMessagingSecureSocket ()
+
+@property(nonatomic, readwrite, assign) FIRMessagingSecureSocketState state;
+
+@end
+
+@interface FIRMessagingClientTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) FIRMessagingClient *client;
+@property(nonatomic, readwrite, strong) id mockClient;
+@property(nonatomic, readwrite, strong) id mockReachability;
+@property(nonatomic, readwrite, strong) id mockRmqManager;
+@property(nonatomic, readwrite, strong) id mockClientDelegate;
+@property(nonatomic, readwrite, strong) id mockDataMessageManager;
+@property(nonatomic, readwrite, strong) id mockRegistrar;
+
+// argument callback blocks
+@property(nonatomic, readwrite, copy) FIRMessagingConnectCompletionHandler connectCompletion;
+@property(nonatomic, readwrite, copy) FIRMessagingTopicOperationCompletion subscribeCompletion;
+
+@end
+
+@implementation FIRMessagingClientTest
+
+- (void)setUp {
+  [super setUp];
+  _mockClientDelegate =
+      OCMStrictProtocolMock(@protocol(FIRMessagingClientDelegate));
+  _mockReachability = OCMClassMock([FIRReachabilityChecker class]);
+  _mockRmqManager = OCMClassMock([FIRMessagingRmqManager class]);
+  _client = [[FIRMessagingClient alloc] initWithDelegate:_mockClientDelegate
+                                   reachability:_mockReachability
+                                    rmq2Manager:_mockRmqManager];
+  _mockClient = OCMPartialMock(_client);
+  _mockRegistrar = OCMPartialMock([_client registrar]);
+  [_mockClient setRegistrar:_mockRegistrar];
+  _mockDataMessageManager = OCMClassMock([FIRMessagingDataMessageManager class]);
+  [_mockClient setDataMessageManager:_mockDataMessageManager];
+}
+
+- (void)tearDown {
+  // remove all handlers
+  [self tearDownMocksAndHandlers];
+  // Mock all sockets to disconnect in a nice way
+  [[[(id)self.client.connection.socket stub] andDo:^(NSInvocation *invocation) {
+      self.client.connection.socket.state = kFIRMessagingSecureSocketClosed;
+  }] disconnect];
+
+  [self.client teardown];
+  [super tearDown];
+}
+
+- (void)tearDownMocksAndHandlers {
+  self.connectCompletion = nil;
+  self.subscribeCompletion = nil;
+}
+
+- (void)setupConnectionWithFakeLoginResult:(BOOL)loginResult
+                          heartbeatTimeout:(NSTimeInterval)heartbeatTimeout {
+  [self setupFakeConnectionWithClass:[FIRMessagingFakeConnection class]
+          withSetupCompletionHandler:^(FIRMessagingConnection *connection) {
+              FIRMessagingFakeConnection *fakeConnection = (FIRMessagingFakeConnection *)connection;
+              fakeConnection.shouldFakeSuccessLogin = loginResult;
+              fakeConnection.fakeConnectionTimeout = heartbeatTimeout;
+          }];
+}
+
+- (void)testSetupConnection {
+  XCTAssertNil(self.client.connection);
+  [self.client setupConnection];
+  XCTAssertNotNil(self.client.connection);
+  XCTAssertNotNil(self.client.connection.delegate);
+}
+
+- (void)testConnectSuccess_withCachedFcmDefaults {
+  [self addFIRMessagingPreferenceKeysToUserDefaults];
+
+  // login request should be successful
+  [self setupConnectionWithFakeLoginResult:YES heartbeatTimeout:1.0];
+
+  XCTestExpectation *setupConnection = [self
+      expectationWithDescription:@"Fcm should successfully setup a connection"];
+
+  [self.client connectWithHandler:^(NSError *error) {
+      XCTAssertNil(error);
+      [setupConnection fulfill];
+  }];
+
+  [self waitForExpectationsWithTimeout:1.0 handler:^(NSError *error) {
+      XCTAssertNil(error);
+  }];
+}
+
+- (void)testsConnectWithNoNetworkError_withCachedFcmDefaults {
+  // connection timeout interval is 1s
+  [[[self.mockClient stub] andReturnValue:@(1)] connectionTimeoutInterval];
+  [self addFIRMessagingPreferenceKeysToUserDefaults];
+
+  [self setupFakeConnectionWithClass:[FIRMessagingFakeFailConnection class]
+          withSetupCompletionHandler:^(FIRMessagingConnection *connection) {
+              FIRMessagingFakeFailConnection *fakeConnection = (FIRMessagingFakeFailConnection *)connection;
+              fakeConnection.shouldFakeSuccessLogin = NO;
+              // should fail only once
+              fakeConnection.failCount = 1;
+          }];
+
+  XCTestExpectation *connectExpectation = [self
+      expectationWithDescription:@"Should retry connection if once failed"];
+  [self.client connectWithHandler:^(NSError *error) {
+      XCTAssertNotNil(error);
+      XCTAssertEqual(kFIRMessagingErrorCodeNetwork, error.code);
+      [connectExpectation fulfill];
+  }];
+
+  [self waitForExpectationsWithTimeout:10.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+}
+
+- (void)testConnectSuccessOnSecondTry_withCachedFcmDefaults {
+  // connection timeout interval is 1s
+  [[[self.mockClient stub] andReturnValue:@(1)] connectionTimeoutInterval];
+  [self addFIRMessagingPreferenceKeysToUserDefaults];
+
+  // the network is available
+  [[[self.mockReachability stub]
+      andReturnValue:@(kFIRReachabilityViaWifi)] reachabilityStatus];
+
+  [self setupFakeConnectionWithClass:[FIRMessagingFakeFailConnection class]
+          withSetupCompletionHandler:^(FIRMessagingConnection *connection) {
+              FIRMessagingFakeFailConnection *fakeConnection = (FIRMessagingFakeFailConnection *)connection;
+              fakeConnection.shouldFakeSuccessLogin = NO;
+              // should fail only once
+              fakeConnection.failCount = 1;
+          }];
+
+  XCTestExpectation *connectExpectation = [self
+      expectationWithDescription:@"Should retry connection if once failed"];
+  [self.client connectWithHandler:^(NSError *error) {
+      XCTAssertNil(error);
+      [connectExpectation fulfill];
+  }];
+
+  [self waitForExpectationsWithTimeout:10.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                                   XCTAssertTrue(
+                                       [self.client isConnectionActive]);
+                               }];
+}
+
+- (void)testDisconnectAfterConnect {
+  // setup the connection
+  [self addFIRMessagingPreferenceKeysToUserDefaults];
+
+  // login request should be successful
+  // Connection should not timeout because of heartbeat failure. Therefore set heartbeatTimeout
+  // to a large value.
+  [self setupConnectionWithFakeLoginResult:YES heartbeatTimeout:100.0];
+
+  [[[self.mockClient stub] andReturnValue:@(1)] connectionTimeoutInterval];
+
+  // the network is available
+  [[[self.mockReachability stub]
+      andReturnValue:@(kFIRReachabilityViaWifi)] reachabilityStatus];
+
+  XCTestExpectation *setupConnection =
+      [self expectationWithDescription:@"Fcm should successfully setup a connection"];
+
+  __block int timesConnected = 0;
+  FIRMessagingConnectCompletionHandler handler = ^(NSError *error) {
+      XCTAssertNil(error);
+      timesConnected++;
+      if (timesConnected == 1) {
+        [setupConnection fulfill];
+        // disconnect the connection after some time
+        FIRMessagingFakeConnection *fakeConnection = (FIRMessagingFakeConnection *)[self.mockClient connection];
+        dispatch_time_t time = dispatch_time(DISPATCH_TIME_NOW, (0.2 * NSEC_PER_SEC));
+        dispatch_after(time, dispatch_get_main_queue(), ^{
+          // disconnect now
+          [(FIRMessagingFakeConnection *)fakeConnection mockSocketDisconnect];
+          [(FIRMessagingFakeConnection *)fakeConnection disconnectNow];
+        });
+      } else {
+        XCTFail(@"Fcm should only connect at max 2 times");
+      }
+  };
+  [self.mockClient connectWithHandler:handler];
+
+  // reconnect after disconnect
+  XCTAssertTrue(self.client.isConnectionActive);
+
+  [self waitForExpectationsWithTimeout:10.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                                   XCTAssertNotEqual(self.client.lastDisconnectedTimestamp, 0);
+                                   XCTAssertTrue(self.client.isConnectionActive);
+                               }];
+}
+
+#pragma mark - Private Helpers
+
+- (void)setupFakeConnectionWithClass:(Class)connectionClass
+          withSetupCompletionHandler:(void (^)(FIRMessagingConnection *))handler {
+  [[[self.mockClient stub] andDo:^(NSInvocation *invocation) {
+      self.client.connection =
+          [[connectionClass alloc] initWithAuthID:kDeviceAuthId
+                                            token:kSecretToken
+                                             host:[FIRMessagingFakeConnection fakeHost]
+                                             port:[FIRMessagingFakeConnection fakePort]
+                                          runLoop:[NSRunLoop mainRunLoop]
+                                      rmq2Manager:self.mockRmqManager
+                                       fcmManager:self.mockDataMessageManager];
+      self.client.connection.delegate = self.client;
+      handler(self.client.connection);
+  }] setupConnection];
+}
+
+- (void)addFIRMessagingPreferenceKeysToUserDefaults {
+  id mockCheckinService = OCMClassMock([FIRMessagingCheckinService class]);
+  [[[mockCheckinService stub] andReturn:kDeviceAuthId] deviceAuthID];
+  [[[mockCheckinService stub] andReturn:kSecretToken] secretToken];
+  [[[mockCheckinService stub] andReturnValue:@YES] hasValidCheckinInfo];
+
+  [[[self.mockRegistrar stub] andReturn:mockCheckinService] checkinService];
+}
+
+@end
+
diff --git a/Example/Messaging/Tests/FIRMessagingCodedInputStreamTest.m b/Example/Messaging/Tests/FIRMessagingCodedInputStreamTest.m
new file mode 100644
index 0000000..7cc2d97
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingCodedInputStreamTest.m
@@ -0,0 +1,116 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import "FIRMessagingCodedInputStream.h"
+
+@interface FIRMessagingCodedInputStreamTest : XCTestCase
+@end
+
+@implementation FIRMessagingCodedInputStreamTest
+
+- (void)testReadingSmallDataStream {
+  FIRMessagingCodedInputStream *stream =
+      [[FIRMessagingCodedInputStream alloc] initWithData:[[self class] sampleData1]];
+  int8_t actualTag = 2;
+  int8_t tag;
+  XCTAssertTrue([stream readTag:&tag]);
+  XCTAssertEqual(actualTag, tag);
+
+  // test length
+  int32_t actualLength = 4;
+  int32_t length;
+  XCTAssertTrue([stream readLength:&length]);
+  XCTAssertEqual(actualLength, length);
+
+  NSData *actualData = [[self class] packetDataForSampleData1];
+  NSData *data = [stream readDataWithLength:length];
+  XCTAssertTrue([actualData isEqualToData:data]);
+}
+
+- (void)testReadingLargeDataStream {
+  FIRMessagingCodedInputStream *stream =
+      [[FIRMessagingCodedInputStream alloc] initWithData:[[self class] sampleData2]];
+  int8_t actualTag = 5;
+  int8_t tag;
+  XCTAssertTrue([stream readTag:&tag]);
+  XCTAssertEqual(actualTag, tag);
+
+  int32_t actualLength = 257;
+  int32_t length;
+  XCTAssertTrue([stream readLength:&length]);
+  XCTAssertEqual(actualLength, length);
+
+  NSData *actualData = [[self class] packetDataForSampleData2];
+  NSData *data = [stream readDataWithLength:length];
+  XCTAssertTrue([actualData isEqualToData:data]);
+}
+
+- (void)testReadingInvalidDataStream {
+  FIRMessagingCodedInputStream *stream =
+      [[FIRMessagingCodedInputStream alloc] initWithData:[[self class] invalidData]];
+  int8_t actualTag = 7;
+  int8_t tag;
+  XCTAssertTrue([stream readTag:&tag]);
+  XCTAssertEqual(actualTag, tag);
+
+  int32_t actualLength = 2;
+  int32_t length;
+  XCTAssertTrue([stream readLength:&length]);
+  XCTAssertEqual(actualLength, length);
+
+  XCTAssertNil([stream readDataWithLength:length]);
+}
+
++ (NSData *)sampleData1 {
+  // tag = 2,
+  // length = 4,
+  // data = integer 255
+  const char data[] = { 0x02, 0x04, 0x80, 0x00, 0x00, 0xff };
+  return [NSData dataWithBytes:data length:6];
+}
+
++ (NSData *)packetDataForSampleData1 {
+  const char data[] = { 0x80, 0x00, 0x00, 0xff };
+  return [NSData dataWithBytes:data length:4];
+}
+
++ (NSData *)sampleData2 {
+  // test reading varint properly
+  // tag = 5,
+  // length = 257,
+  // data = length 257
+  const char tagAndLength[] = { 0x05, 0x81, 0x02 };
+  NSMutableData *data = [NSMutableData dataWithBytes:tagAndLength length:3];
+  [data appendData:[self packetDataForSampleData2]];
+  return data;
+}
+
++ (NSData *)packetDataForSampleData2 {
+  char packetData[257] = { 0xff, 0xff, 0xff };
+  return [NSData dataWithBytes:packetData length:257];
+}
+
++ (NSData *)invalidData {
+  // tag = 7,
+  // length = 2,
+  // data = (length 1)
+  const char data[] = { 0x07, 0x02, 0xff };
+  return [NSData dataWithBytes:data length:3];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingConnectionTest.m b/Example/Messaging/Tests/FIRMessagingConnectionTest.m
new file mode 100644
index 0000000..47b29d2
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingConnectionTest.m
@@ -0,0 +1,480 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+#import <GoogleToolboxForMac/GTMDefines.h>
+
+#import "FIRMessagingClient.h"
+#import "FIRMessagingConnection.h"
+#import "FIRMessagingDataMessageManager.h"
+#import "FIRMessagingFakeConnection.h"
+#import "FIRMessagingRmqManager.h"
+#import "FIRMessagingSecureSocket.h"
+#import "FIRMessagingUtilities.h"
+
+static NSString *const kDeviceAuthId = @"123456";
+static NSString *const kSecretToken = @"56789";
+
+// used to verify if we are sending in the right proto or not.
+// set it to negative value to disable this check
+static FIRMessagingProtoTag currentProtoSendTag;
+
+@interface FIRMessagingSecureSocket ()
+
+@property(nonatomic, readwrite, assign) FIRMessagingSecureSocketState state;
+
+@end
+
+@interface FIRMessagingSecureSocket (test_FIRMessagingConnection)
+
+- (void)_successconnectToHost:(NSString *)host
+                         port:(NSUInteger)port
+                    onRunLoop:(NSRunLoop *)runLoop;
+- (void)_fakeSuccessfulSocketConnect;
+
+@end
+
+@implementation FIRMessagingSecureSocket (test_FIRMessagingConnection)
+
+- (void)_successconnectToHost:(NSString *)host
+                         port:(NSUInteger)port
+                    onRunLoop:(NSRunLoop *)runLoop {
+  // created ports, opened streams
+  // invoke callback async
+  [self _fakeSuccessfulSocketConnect];
+}
+
+- (void)_fakeSuccessfulSocketConnect {
+  self.state = kFIRMessagingSecureSocketOpen;
+  [self.delegate secureSocketDidConnect:self];
+}
+
+@end
+
+// make sure these are defined in FIRMessagingConnection
+@interface FIRMessagingConnection () <FIRMessagingSecureSocketDelegate>
+
+@property(nonatomic, readwrite, assign) int64_t lastLoginServerTimestamp;
+@property(nonatomic, readwrite, assign) int lastStreamIdAcked;
+@property(nonatomic, readwrite, assign) int inStreamId;
+@property(nonatomic, readwrite, assign) int outStreamId;
+
+@property(nonatomic, readwrite, strong) FIRMessagingSecureSocket *socket;
+
+@property(nonatomic, readwrite, strong) NSMutableArray *unackedS2dIds;
+@property(nonatomic, readwrite, strong) NSMutableDictionary *ackedS2dMap;
+@property(nonatomic, readwrite, strong) NSMutableArray *d2sInfos;
+
+- (void)setupConnectionSocket;
+- (void)connectToSocket:(FIRMessagingSecureSocket *)socket;
+- (NSTimeInterval)connectionTimeoutInterval;
+- (void)sendHeartbeatPing;
+
+@end
+
+
+@interface FIRMessagingConnectionTest : XCTestCase
+
+@property(nonatomic, readwrite, assign) BOOL didSuccessfullySendData;
+
+@property(nonatomic, readwrite, strong) NSUserDefaults *userDefaults;
+@property(nonatomic, readwrite, strong) FIRMessagingConnection *fakeConnection;
+@property(nonatomic, readwrite, strong) id mockClient;
+@property(nonatomic, readwrite, strong) id mockConnection;
+@property(nonatomic, readwrite, strong) id mockRmq;
+@property(nonatomic, readwrite, strong) id mockDataMessageManager;
+
+@end
+
+@implementation FIRMessagingConnectionTest
+
+- (void)setUp {
+  [super setUp];
+  _userDefaults = [[NSUserDefaults alloc] init];
+  _mockRmq = OCMClassMock([FIRMessagingRmqManager class]);
+  _mockDataMessageManager = OCMClassMock([FIRMessagingDataMessageManager class]);
+  // fake connection is only used to simulate the socket behavior
+  _fakeConnection = [[FIRMessagingFakeConnection alloc] initWithAuthID:kDeviceAuthId
+                                                        token:kSecretToken
+                                                         host:[FIRMessagingFakeConnection fakeHost]
+                                                         port:[FIRMessagingFakeConnection fakePort]
+                                                      runLoop:[NSRunLoop currentRunLoop]
+                                                  rmq2Manager:_mockRmq
+                                                   fcmManager:_mockDataMessageManager];
+
+  _mockClient = OCMClassMock([FIRMessagingClient class]);
+  _fakeConnection.delegate = _mockClient;
+  _mockConnection = OCMPartialMock(_fakeConnection);
+  _didSuccessfullySendData = NO;
+}
+
+- (void)tearDown {
+  [self.fakeConnection teardown];
+  [super tearDown];
+}
+
+- (void)testInitialConnectionNotConnected {
+  XCTAssertEqual(kFIRMessagingConnectionNotConnected, [self.fakeConnection state]);
+}
+
+- (void)testSuccessfulSocketConnection {
+  [self.fakeConnection signIn];
+
+  // should be connected now
+  XCTAssertEqual(kFIRMessagingConnectionConnected, self.fakeConnection.state);
+  XCTAssertEqual(0, self.fakeConnection.lastStreamIdAcked);
+  XCTAssertEqual(0, self.fakeConnection.inStreamId);
+  XCTAssertEqual(0, self.fakeConnection.ackedS2dMap.count);
+  XCTAssertEqual(0, self.fakeConnection.unackedS2dIds.count);
+
+  [self stubSocketDisconnect:self.fakeConnection.socket];
+}
+
+- (void)testSignInAndThenSignOut {
+  [self.fakeConnection signIn];
+  [self stubSocketDisconnect:self.fakeConnection.socket];
+  [self.fakeConnection signOut];
+  XCTAssertEqual(kFIRMessagingSecureSocketClosed, self.fakeConnection.socket.state);
+}
+
+- (void)testSuccessfulSignIn {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+  XCTAssertEqual(self.fakeConnection.state, kFIRMessagingConnectionSignedIn);
+  XCTAssertEqual(self.fakeConnection.outStreamId, 2);
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testSignOut_whenSignedIn {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+
+  // should be signed in now
+  id mockSocket = self.fakeConnection.socket;
+  [self.fakeConnection signOut];
+  XCTAssertEqual(self.fakeConnection.state, kFIRMessagingConnectionNotConnected);
+  XCTAssertEqual(self.fakeConnection.outStreamId, 3);
+  XCTAssertNil([(FIRMessagingSecureSocket *)mockSocket delegate]);
+  XCTAssertTrue(self.didSuccessfullySendData);
+  OCMVerify([mockSocket sendData:[OCMArg any]
+                         withTag:kFIRMessagingProtoTagClose
+                           rmqId:[OCMArg isNil]]);
+}
+
+- (void)testReceiveCloseProto {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+
+  id mockSocket = self.fakeConnection.socket;
+  GtalkClose *close = [[GtalkClose alloc] init];
+  [self.fakeConnection secureSocket:mockSocket
+                     didReceiveData:[close data]
+                            withTag:kFIRMessagingProtoTagClose];
+  XCTAssertEqual(self.fakeConnection.state, kFIRMessagingConnectionNotConnected);
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testLoginRequest {
+  XCTAssertEqual(kFIRMessagingConnectionNotConnected, [self.fakeConnection state]);
+  [self.fakeConnection setupConnectionSocket];
+
+  id socketMock = OCMPartialMock(self.fakeConnection.socket);
+  self.fakeConnection.socket = socketMock;
+
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        [socketMock _fakeSuccessfulSocketConnect];
+      }]
+      connectToHost:[FIRMessagingFakeConnection fakeHost]
+               port:[FIRMessagingFakeConnection fakePort]
+          onRunLoop:[OCMArg any]];
+
+  [[[socketMock stub] andCall:@selector(_sendData:withTag:rmqId:) onObject:self]
+      // do nothing
+      sendData:[OCMArg any]
+       withTag:kFIRMessagingProtoTagLoginRequest
+         rmqId:[OCMArg isNil]];
+
+  // swizzle disconnect socket
+  OCMVerify([[[socketMock stub] andCall:@selector(_disconnectSocket)
+                               onObject:self] disconnect]);
+
+  currentProtoSendTag = kFIRMessagingProtoTagLoginRequest;
+  // send login request
+  [self.fakeConnection connectToSocket:socketMock];
+
+  // verify login request sent
+  XCTAssertEqual(1, self.fakeConnection.outStreamId);
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testLoginRequest_withPendingMessagesInRmq {
+  // TODO: add fake messages to rmq and test login request with them
+}
+
+- (void)testLoginRequest_withSuccessfulResponse {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+
+  OCMVerify([self.mockClient didLoginWithConnection:[OCMArg isEqual:self.fakeConnection]]);
+
+  // should send a heartbeat ping too
+  XCTAssertEqual(self.fakeConnection.outStreamId, 2);
+  // update for the received login response proto
+  XCTAssertEqual(self.fakeConnection.inStreamId, 1);
+  // did send data during login
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testConnectionTimeout {
+  XCTAssertEqual(kFIRMessagingConnectionNotConnected, [self.fakeConnection state]);
+
+  [self.fakeConnection setupConnectionSocket];
+
+  id socketMock = OCMPartialMock(self.fakeConnection.socket);
+  self.fakeConnection.socket = socketMock;
+
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        [socketMock _fakeSuccessfulSocketConnect];
+      }]
+      connectToHost:[FIRMessagingFakeConnection fakeHost]
+               port:[FIRMessagingFakeConnection fakePort]
+          onRunLoop:[OCMArg any]];
+
+  [self.fakeConnection connectToSocket:socketMock];
+  XCTAssertEqual(self.fakeConnection.state, kFIRMessagingConnectionConnected);
+
+  GtalkLoginResponse *response = [[GtalkLoginResponse alloc] init];
+  [response setId_p:@""];
+
+  // connection timeout has been scheduled
+  // should disconnect since we wait for more time
+  XCTestExpectation *disconnectExpectation =
+      [self expectationWithDescription:
+          @"FCM connection should timeout without receiving "
+          @"any data for a timeout interval"];
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        [self _disconnectSocket];
+        [disconnectExpectation fulfill];
+      }] disconnect];
+
+  // simulate connection receiving login response
+  [self.fakeConnection secureSocket:socketMock
+                     didReceiveData:[response data]
+                            withTag:kFIRMessagingProtoTagLoginResponse];
+
+  [self waitForExpectationsWithTimeout:2.0
+                               handler:^(NSError *error) {
+                                 XCTAssertNil(error);
+                               }];
+
+  [socketMock verify];
+  XCTAssertEqual(self.fakeConnection.state, kFIRMessagingConnectionNotConnected);
+}
+
+- (void)testDataMessageReceive {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+  GtalkDataMessageStanza *stanza = [[GtalkDataMessageStanza alloc] init];
+  [stanza setCategory:@"special"];
+  [stanza setFrom:@"xyz"];
+  [self.fakeConnection secureSocket:self.fakeConnection.socket
+                     didReceiveData:[stanza data]
+                            withTag:kFIRMessagingProtoTagDataMessageStanza];
+
+  OCMVerify([self.mockClient connectionDidRecieveMessage:[OCMArg checkWithBlock:^BOOL(id obj) {
+      GtalkDataMessageStanza *message = (GtalkDataMessageStanza *)obj;
+      return [[message category] isEqual:@"special"] && [[message from] isEqual:@"xyz"];
+  }]]);
+  // did send data while login
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testDataMessageReceiveWithInvalidTag {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+  GtalkDataMessageStanza *stanza = [[GtalkDataMessageStanza alloc] init];
+  BOOL didCauseException = NO;
+  @try {
+    [self.fakeConnection secureSocket:self.fakeConnection.socket
+                       didReceiveData:[stanza data]
+                              withTag:kFIRMessagingProtoTagInvalid];
+  } @catch (NSException *exception) {
+    didCauseException = YES;
+  } @finally {
+  }
+  XCTAssertFalse(didCauseException);
+}
+
+- (void)testDataMessageReceiveWithTagThatDoesntEquateToClass {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+  GtalkDataMessageStanza *stanza = [[GtalkDataMessageStanza alloc] init];
+  BOOL didCauseException = NO;
+  int8_t tagWhichDoesntEquateToClass = INT8_MAX;
+  @try {
+    [self.fakeConnection secureSocket:self.fakeConnection.socket
+                       didReceiveData:[stanza data]
+                              withTag:tagWhichDoesntEquateToClass];
+  } @catch (NSException *exception) {
+    didCauseException = YES;
+  } @finally {
+  }
+  XCTAssertFalse(didCauseException);
+}
+
+- (void)testHeartbeatSend {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection]; // outstreamId should be 2
+  XCTAssertEqual(self.fakeConnection.outStreamId, 2);
+  [self.fakeConnection sendHeartbeatPing];
+  id mockSocket = self.fakeConnection.socket;
+  OCMVerify([mockSocket sendData:[OCMArg any]
+                         withTag:kFIRMessagingProtoTagHeartbeatPing
+                           rmqId:[OCMArg isNil]]);
+  XCTAssertEqual(self.fakeConnection.outStreamId, 3);
+  // did send data
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+- (void)testHeartbeatReceived {
+  [self setupSuccessfulLoginRequestWithConnection:self.fakeConnection];
+  XCTAssertEqual(self.fakeConnection.outStreamId, 2);
+  GtalkHeartbeatPing *ping = [[GtalkHeartbeatPing alloc] init];
+  [self.fakeConnection secureSocket:self.fakeConnection.socket
+                     didReceiveData:[ping data]
+                            withTag:kFIRMessagingProtoTagHeartbeatPing];
+  XCTAssertEqual(self.fakeConnection.inStreamId, 2);
+  id mockSocket = self.fakeConnection.socket;
+  OCMVerify([mockSocket sendData:[OCMArg any]
+                         withTag:kFIRMessagingProtoTagHeartbeatAck
+                           rmqId:[OCMArg isNil]]);
+  XCTAssertEqual(self.fakeConnection.outStreamId, 3);
+  // did send data
+  XCTAssertTrue(self.didSuccessfullySendData);
+}
+
+// TODO: Add tests for Selective/Stream ACK's
+
+#pragma mark - Stubs
+
+- (void)_disconnectSocket {
+  self.fakeConnection.socket.state = kFIRMessagingSecureSocketClosed;
+}
+
+- (void)_sendData:(NSData *)data withTag:(int8_t)tag rmqId:(NSString *)rmqId {
+  _GTMDevLog(@"FIRMessaging Socket: Send data with Tag: %d rmq: %@", tag, rmqId);
+  if (currentProtoSendTag > 0) {
+    XCTAssertEqual(tag, currentProtoSendTag);
+  }
+  self.didSuccessfullySendData = YES;
+}
+
+#pragma mark - Private Helpers
+
+/**
+ * Stub socket disconnect to prevent spurious assert. Since we mock the socket object being
+ * used by the connection, while we teardown the client we also disconnect the socket to tear
+ * it down. Since we are using mock sockets we need to stub the `disconnect` to prevent some
+ * assertions from taking place.
+ * The `_disconectSocket` has the gist of the actual socket disconnect without any assertions.
+ */
+- (void)stubSocketDisconnect:(id)mockSocket {
+  [[[mockSocket stub] andCall:@selector(_disconnectSocket)
+                     onObject:self] disconnect];
+
+  [mockSocket verify];
+}
+
+- (void)mockSuccessfulSignIn {
+  XCTAssertEqual(kFIRMessagingConnectionNotConnected, [self.fakeConnection state]);
+  [self.fakeConnection setupConnectionSocket];
+
+  id socketMock = OCMPartialMock(self.fakeConnection.socket);
+  self.fakeConnection.socket = socketMock;
+
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        [socketMock _fakeSuccessfulSocketConnect];
+      }]
+      connectToHost:[FIRMessagingFakeConnection fakeHost]
+               port:[FIRMessagingFakeConnection fakePort]
+          onRunLoop:[OCMArg any]];
+
+  [[[socketMock stub] andCall:@selector(_sendData:withTag:rmqId:) onObject:self]
+      // do nothing
+      sendData:[OCMArg any]
+       withTag:kFIRMessagingProtoTagLoginRequest
+         rmqId:[OCMArg isNil]];
+
+  // send login request
+  currentProtoSendTag = kFIRMessagingProtoTagLoginRequest;
+  [self.fakeConnection connectToSocket:socketMock];
+
+  GtalkLoginResponse *response = [[GtalkLoginResponse alloc] init];
+  [response setId_p:@""];
+
+  // simulate connection receiving login response
+  [self.fakeConnection secureSocket:socketMock
+                     didReceiveData:[response data]
+                            withTag:kFIRMessagingProtoTagLoginResponse];
+
+  OCMVerify([self.mockClient didLoginWithConnection:[OCMArg isEqual:self.fakeConnection]]);
+
+  // should receive data
+  XCTAssertTrue(self.didSuccessfullySendData);
+  // should send a heartbeat ping too
+  XCTAssertEqual(self.fakeConnection.outStreamId, 2);
+  // update for the received login response proto
+  XCTAssertEqual(self.fakeConnection.inStreamId, 1);
+}
+
+- (void)setupSuccessfulLoginRequestWithConnection:(FIRMessagingConnection *)fakeConnection {
+  [fakeConnection setupConnectionSocket];
+
+  id socketMock = OCMPartialMock(fakeConnection.socket);
+  fakeConnection.socket = socketMock;
+
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        [socketMock _fakeSuccessfulSocketConnect];
+      }]
+      connectToHost:[FIRMessagingFakeConnection fakeHost]
+               port:[FIRMessagingFakeConnection fakePort]
+          onRunLoop:[OCMArg any]];
+
+  [[[socketMock stub] andCall:@selector(_sendData:withTag:rmqId:) onObject:self]
+       // do nothing
+       sendData:[OCMArg any]
+        withTag:kFIRMessagingProtoTagLoginRequest
+          rmqId:[OCMArg isNil]];
+
+  // swizzle disconnect socket
+  [[[socketMock stub] andCall:@selector(_disconnectSocket)
+                     onObject:self] disconnect];
+
+  // send login request
+  currentProtoSendTag = kFIRMessagingProtoTagLoginRequest;
+  [fakeConnection connectToSocket:socketMock];
+
+  GtalkLoginResponse *response = [[GtalkLoginResponse alloc] init];
+  [response setId_p:@""];
+
+  // simulate connection receiving login response
+  [fakeConnection secureSocket:socketMock
+                didReceiveData:[response data]
+                       withTag:kFIRMessagingProtoTagLoginResponse];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingContextManagerServiceTest.m b/Example/Messaging/Tests/FIRMessagingContextManagerServiceTest.m
new file mode 100644
index 0000000..cb48e7f
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingContextManagerServiceTest.m
@@ -0,0 +1,183 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <XCTest/XCTest.h>
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessagingContextManagerService.h"
+
+@interface FIRMessagingContextManagerServiceTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) NSDateFormatter *dateFormatter;
+@property(nonatomic, readwrite, strong) NSMutableArray *scheduledLocalNotifications;
+
+@end
+
+@implementation FIRMessagingContextManagerServiceTest
+
+- (void)setUp {
+  [super setUp];
+  self.dateFormatter = [[NSDateFormatter alloc] init];
+  [self.dateFormatter setDateFormat:@"yyyy-MM-dd HH:mm:ss"];
+  self.scheduledLocalNotifications = [NSMutableArray array];
+  [self mockSchedulingLocalNotifications];
+}
+
+- (void)tearDown {
+  [super tearDown];
+}
+
+/**
+ *  Test invalid context manager message, missing lt_start string.
+ */
+- (void)testInvalidContextManagerMessage_missingStartTime {
+  NSDictionary *message = @{
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([FIRMessagingContextManagerService isContextManagerMessage:message]);
+}
+
+/**
+ *  Test valid context manager message.
+ */
+- (void)testValidContextManagerMessage {
+  NSDictionary *message = @{
+    kFIRMessagingContextManagerLocalTimeStart: @"2015-12-12 00:00:00",
+    @"hello" : @"world",
+  };
+  XCTAssertTrue([FIRMessagingContextManagerService isContextManagerMessage:message]);
+}
+
+// TODO: Enable these tests. They fail because we cannot schedule local
+// notifications on OSX without permission. It's better to mock AppDelegate's
+// scheduleLocalNotification to mock scheduling behavior.
+
+/**
+ *  Context Manager message with future start date should be successfully scheduled.
+ */
+- (void)testMessageWithFutureStartTime {
+  NSString *messageIdentifier = @"fcm-cm-test1";
+  NSString *startTimeString = @"2020-01-12 12:00:00";  // way into the future
+  NSDictionary *message = @{
+    kFIRMessagingContextManagerLocalTimeStart: startTimeString,
+    kFIRMessagingContextManagerBodyKey : @"Hello world!",
+    @"id": messageIdentifier,
+    @"hello" : @"world"
+  };
+
+  XCTAssertTrue([FIRMessagingContextManagerService handleContextManagerMessage:message]);
+
+  XCTAssertEqual(self.scheduledLocalNotifications.count, 1);
+  UILocalNotification *notification = [self.scheduledLocalNotifications firstObject];
+  NSDate *date = [self.dateFormatter dateFromString:startTimeString];
+  XCTAssertEqual([notification.fireDate compare:date], NSOrderedSame);
+}
+
+/**
+ *  Context Manager message with past end date should not be scheduled.
+ */
+- (void)testMessageWithPastEndTime {
+  NSString *messageIdentifier = @"fcm-cm-test1";
+  NSString *startTimeString = @"2010-01-12 12:00:00";  // way into the past
+  NSString *endTimeString = @"2011-01-12 12:00:00";  // way into the past
+  NSDictionary *message = @{
+    kFIRMessagingContextManagerLocalTimeStart: startTimeString,
+    kFIRMessagingContextManagerLocalTimeEnd : endTimeString,
+    kFIRMessagingContextManagerBodyKey : @"Hello world!",
+    @"id": messageIdentifier,
+    @"hello" : @"world"
+  };
+
+  XCTAssertTrue([FIRMessagingContextManagerService handleContextManagerMessage:message]);
+  XCTAssertEqual(self.scheduledLocalNotifications.count, 0);
+}
+
+/**
+ *  Context Manager message with past start and future end date should be successfully
+ *  scheduled.
+ */
+- (void)testMessageWithPastStartAndFutureEndTime {
+  NSString *messageIdentifier = @"fcm-cm-test1";
+  NSDate *startDate = [NSDate dateWithTimeIntervalSinceNow:-1000];  // past
+  NSDate *endDate = [NSDate dateWithTimeIntervalSinceNow:1000];  // future
+  NSString *startTimeString = [self.dateFormatter stringFromDate:startDate];
+  NSString *endTimeString = [self.dateFormatter stringFromDate:endDate];
+
+  NSDictionary *message = @{
+    kFIRMessagingContextManagerLocalTimeStart : startTimeString,
+    kFIRMessagingContextManagerLocalTimeEnd : endTimeString,
+    kFIRMessagingContextManagerBodyKey : @"Hello world!",
+    @"id": messageIdentifier,
+    @"hello" : @"world"
+  };
+
+  XCTAssertTrue([FIRMessagingContextManagerService handleContextManagerMessage:message]);
+
+  XCTAssertEqual(self.scheduledLocalNotifications.count, 1);
+  UILocalNotification *notification = [self.scheduledLocalNotifications firstObject];
+  // schedule notification after start date
+  XCTAssertEqual([notification.fireDate compare:startDate], NSOrderedDescending);
+  // schedule notification after end date
+  XCTAssertEqual([notification.fireDate compare:endDate], NSOrderedAscending);
+}
+
+/**
+ *  Test correctly parsing user data in local notifications.
+ */
+- (void)testTimedNotificationsUserInfo {
+  NSString *messageIdentifierKey = @"message.id";
+  NSString *messageIdentifier = @"fcm-cm-test1";
+  NSString *startTimeString = @"2020-01-12 12:00:00";  // way into the future
+
+  NSString *customDataKey = @"hello";
+  NSString *customData = @"world";
+  NSDictionary *message = @{
+    kFIRMessagingContextManagerLocalTimeStart : startTimeString,
+    kFIRMessagingContextManagerBodyKey : @"Hello world!",
+    messageIdentifierKey : messageIdentifier,
+    customDataKey : customData,
+  };
+
+  XCTAssertTrue([FIRMessagingContextManagerService handleContextManagerMessage:message]);
+
+  XCTAssertEqual(self.scheduledLocalNotifications.count, 1);
+  UILocalNotification *notification = [self.scheduledLocalNotifications firstObject];
+  XCTAssertEqualObjects(notification.userInfo[messageIdentifierKey], messageIdentifier);
+  XCTAssertEqualObjects(notification.userInfo[customDataKey], customData);
+}
+
+#pragma mark - Private Helpers
+
+- (void)mockSchedulingLocalNotifications {
+  id mockApplication = OCMPartialMock([UIApplication sharedApplication]);
+  __block UILocalNotification *notificationToSchedule;
+  [[[mockApplication stub]
+      andDo:^(NSInvocation *invocation) {
+        // Mock scheduling a notification
+        if (notificationToSchedule) {
+          [self.scheduledLocalNotifications addObject:notificationToSchedule];
+        }
+      }] scheduleLocalNotification:[OCMArg checkWithBlock:^BOOL(id obj) {
+        if ([obj isKindOfClass:[UILocalNotification class]]) {
+          notificationToSchedule = obj;
+          return YES;
+        }
+        return NO;
+      }]];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingDataMessageManagerTest.m b/Example/Messaging/Tests/FIRMessagingDataMessageManagerTest.m
new file mode 100644
index 0000000..2b4f407
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingDataMessageManagerTest.m
@@ -0,0 +1,662 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+
+#import "FIRMessaging.h"
+#import "FIRMessagingClient.h"
+#import "FIRMessagingConfig.h"
+#import "FIRMessagingConnection.h"
+#import "FIRMessagingDataMessageManager.h"
+#import "FIRMessagingReceiver.h"
+#import "FIRMessagingRmqManager.h"
+#import "FIRMessagingSyncMessageManager.h"
+#import "FIRMessagingUtilities.h"
+#import "FIRMessaging_Private.h"
+#import "FIRMessagingConstants.h"
+#import "FIRMessagingDefines.h"
+#import "NSError+FIRMessaging.h"
+
+static NSString *const kFIRMessagingUserDefaultsSuite = @"FIRMessagingClientTestUserDefaultsSuite";
+
+static NSString *const kFIRMessagingAppIDToken = @"1234abcdef789";
+
+static NSString *const kMessagePersistentID = @"abcdef123";
+static NSString *const kMessageFrom = @"com.example.gcm";
+static NSString *const kMessageTo = @"123456789";
+static NSString *const kCollapseKey = @"collapse-1";
+static NSString *const kAppDataItemKey = @"hello";
+static NSString *const kAppDataItemValue = @"world";
+static NSString *const kAppDataItemInvalidKey = @"google.hello";
+
+static NSString *const kRmqDatabaseName = @"gcm-dmm-test";
+
+@interface FIRMessagingDataMessageManager()
+
+@property(nonatomic, readwrite, weak) FIRMessagingRmqManager *rmq2Manager;
+
+- (NSString *)categoryForUpstreamMessages;
+
+@end
+
+@interface FIRMessagingDataMessageManagerTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) id mockClient;
+@property(nonatomic, readwrite, strong) id mockRmqManager;
+@property(nonatomic, readwrite, strong) id mockReceiver;
+@property(nonatomic, readwrite, strong) id mockSyncMessageManager;
+@property(nonatomic, readwrite, strong) FIRMessagingDataMessageManager *dataMessageManager;
+@property(nonatomic, readwrite, strong) id mockDataMessageManager;
+
+@end
+
+@implementation FIRMessagingDataMessageManagerTest
+
+- (void)setUp {
+  [super setUp];
+  _mockClient = OCMClassMock([FIRMessagingClient class]);
+  _mockReceiver = OCMClassMock([FIRMessagingReceiver class]);
+  _mockRmqManager = OCMClassMock([FIRMessagingRmqManager class]);
+  _mockSyncMessageManager = OCMClassMock([FIRMessagingSyncMessageManager class]);
+  _dataMessageManager = [[FIRMessagingDataMessageManager alloc]
+        initWithDelegate:_mockReceiver
+                  client:_mockClient
+             rmq2Manager:_mockRmqManager
+      syncMessageManager:_mockSyncMessageManager];
+  [_dataMessageManager refreshDelayedMessages];
+  _mockDataMessageManager = OCMPartialMock(_dataMessageManager);
+}
+
+
+- (void)testSendValidMessage_withNoConnection {
+  // mock no connection initially
+  NSString *messageID = @"1";
+  BOOL mockConnectionActive = NO;
+  [[[self.mockClient stub] andDo:^(NSInvocation *invocation) {
+    NSValue *returnValue = [NSValue valueWithBytes:&mockConnectionActive
+                                          objCType:@encode(BOOL)];
+    [invocation setReturnValue:&returnValue];
+  }] isConnectionActive];
+
+  BOOL(^isValidStanza)(id obj) = ^BOOL(id obj) {
+    if ([obj isKindOfClass:[GtalkDataMessageStanza class]]) {
+      GtalkDataMessageStanza *message = (GtalkDataMessageStanza *)obj;
+      return ([message.id_p isEqualToString:messageID] && [message.to isEqualToString:kMessageTo]);
+    }
+    return NO;
+  };
+  OCMExpect([self.mockReceiver willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                                                   error:[OCMArg isNil]]);
+  [[[self.mockRmqManager stub] andReturnValue:@YES]
+      saveRmqMessage:[OCMArg checkWithBlock:isValidStanza]
+               error:[OCMArg anyObjectRef]];
+
+  // should be logged into the service
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  // try to send messages with no connection should be queued into RMQ
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:-1 delay:0];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+  OCMVerifyAll(self.mockRmqManager);
+}
+
+- (void)testSendValidMessage_withoutCheckinAuthentication {
+  NSString *messageID = @"1";
+  NSMutableDictionary *message = [self standardFIRMessagingMessageWithMessageID:messageID];
+
+  OCMExpect([self.mockReceiver
+      willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                          error:[OCMArg checkWithBlock:^BOOL(id obj) {
+                            if ([obj isKindOfClass:[NSError class]]) {
+                              NSError *error = (NSError *)obj;
+                              return error.code == kFIRMessagingErrorCodeMissingDeviceID;
+                            }
+                            return NO;
+                          }]]);
+
+  // do not log into checkin service
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+}
+
+- (void)testSendInvalidMessage_withNoTo {
+  NSString *messageID = @"1";
+  NSMutableDictionary *message =
+      [FIRMessaging createFIRMessagingMessageWithMessage:@{ kAppDataItemKey : kAppDataItemValue}
+                                                      to:@""
+                                                  withID:messageID
+                                              timeToLive:-1
+                                                   delay:0];
+
+  OCMExpect([self.mockReceiver
+      willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                          error:[OCMArg checkWithBlock:^BOOL(id obj) {
+                            if ([obj isKindOfClass:[NSError class]]) {
+                              NSError *error = (NSError *)obj;
+                              return error.code == kFIRMessagingErrorMissingTo;
+                            }
+                            return NO;
+                          }]]);
+
+  // should be logged into the service
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+}
+
+- (void)testSendInvalidMessage_withSizeExceeded {
+  NSString *messageID = @"1";
+  NSString *veryLargeString = [@"a" stringByPaddingToLength:4 * 1024 // 4kB
+                                                 withString:@"b"
+                                            startingAtIndex:0];
+  NSMutableDictionary *message =
+      [FIRMessaging createFIRMessagingMessageWithMessage:@{ kAppDataItemKey : veryLargeString }
+                                                      to:kMessageTo
+                                                  withID:messageID
+                                              timeToLive:-1
+                                                   delay:0];
+
+  OCMExpect([self.mockReceiver
+      willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                          error:[OCMArg checkWithBlock:^BOOL(id obj) {
+                            if ([obj isKindOfClass:[NSError class]]) {
+                              NSError *error = (NSError *)obj;
+                              return error.code == kFIRMessagingErrorSizeExceeded;
+                            }
+                            return NO;
+                          }]]);
+
+  [self addFakeFIRMessagingRegistrationToken];
+  // should be logged into the service
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+}
+
+// TODO: Add test with rawData exceeding 4KB in size
+
+- (void)testSendValidMessage_withRmqSaveError {
+  NSString *messageID = @"1";
+  NSMutableDictionary *message = [self standardFIRMessagingMessageWithMessageID:messageID];
+  [[[self.mockRmqManager stub] andReturnValue:@NO]
+      saveRmqMessage:[OCMArg any] error:[OCMArg anyObjectRef]];
+
+  OCMExpect([self.mockReceiver
+      willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                          error:[OCMArg checkWithBlock:^BOOL(id obj) {
+                            if ([obj isKindOfClass:[NSError class]]) {
+                              NSError *error = (NSError *)obj;
+                              return error.code == kFIRMessagingErrorSave;
+                            }
+                            return NO;
+                          }]]);
+
+  // should be logged into the service
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+}
+
+- (void)testSendValidMessage_withTTL0 {
+  // simulate a valid connection
+  [[[self.mockClient stub] andReturnValue:@YES] isConnectionActive];
+  NSString *messageID = @"1";
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:0 delay:0];
+
+  BOOL(^isValidStanza)(id obj) = ^BOOL(id obj) {
+    if ([obj isKindOfClass:[GtalkDataMessageStanza class]]) {
+      GtalkDataMessageStanza *stanza = (GtalkDataMessageStanza *)obj;
+      return ([stanza.id_p isEqualToString:messageID] &&
+              [stanza.to isEqualToString:kMessageTo] &&
+              stanza.ttl == 0);
+    }
+    return NO;
+  };
+
+  OCMExpect([self.mockClient sendMessage:[OCMArg checkWithBlock:isValidStanza]]);
+
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockClient);
+}
+
+// TODO: This is failing on simulator 7.1 & 8.2, take this out temporarily
+- (void)XXX_testSendValidMessage_withTTL0AndNoFIRMessagingConnection {
+  // simulate a invalid connection
+  [[[self.mockClient stub] andReturnValue:@NO] isConnectionActive];
+
+  // simulate network reachability
+  FIRMessaging *service = [FIRMessaging messaging];
+  id mockService = OCMPartialMock(service);
+  [[[mockService stub] andReturnValue:@YES] isNetworkAvailable];
+
+  NSString *messageID = @"1";
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:0 delay:0];
+
+
+  BOOL(^isValidStanza)(id obj) = ^BOOL(id obj) {
+    if ([obj isKindOfClass:[GtalkDataMessageStanza class]]) {
+      GtalkDataMessageStanza *stanza = (GtalkDataMessageStanza *)obj;
+      return ([stanza.id_p isEqualToString:messageID] &&
+              [stanza.to isEqualToString:kMessageTo] &&
+              stanza.ttl == 0);
+    }
+    return NO;
+  };
+
+  // should save the message to be sent when we reconnect the next time
+  OCMExpect([self.mockClient sendOnConnectOrDrop:[OCMArg checkWithBlock:isValidStanza]]);
+  // should also try to reconnect immediately
+  OCMExpect([self.mockClient retryConnectionImmediately:[OCMArg isEqual:@YES]]);
+
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockClient);
+}
+
+// TODO: Investigate why this test is flaky
+- (void)xxx_testSendValidMessage_withTTL0AndNoNetwork {
+  // simulate a invalid connection
+  [[[self.mockClient stub] andReturnValue:@NO] isConnectionActive];
+
+  NSString *messageID = @"1";
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:0 delay:0];
+
+
+  // should drop the message since there is no network
+  OCMExpect([self.mockReceiver willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                                                   error:[OCMArg checkWithBlock:^BOOL(id obj) {
+    if ([obj isKindOfClass:[NSError class]]) {
+      NSError *error = (NSError *)obj;
+      return error.code == kFIRMessagingErrorCodeNetwork;
+    }
+    return NO;
+  }]]);
+
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockReceiver);
+}
+
+// TODO: This failed on simulator 7.1 & 8.2, take this out temporarily
+- (void)XXX_testDelayedMessagesBeingResentOnReconnect {
+  static BOOL isConnectionActive = NO;
+  OCMStub([self.mockClient isConnectionActive]).andDo(^(NSInvocation *invocation) {
+    [invocation setReturnValue:&isConnectionActive];
+  });
+
+  // message that lives for 2 seconds
+  NSString *messageID = @"1";
+  int ttl = 2;
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:ttl delay:1];
+
+  __block GtalkDataMessageStanza *firstMessageStanza;
+
+  OCMStub([self.mockRmqManager saveRmqMessage:[OCMArg any]
+                                        error:[OCMArg anyObjectRef]]).andReturn(YES);
+
+  OCMExpect([self.mockReceiver willSendDataMessageWithID:[OCMArg isEqual:messageID]
+                                                   error:[OCMArg isNil]]);
+
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  __block FIRMessagingDataMessageHandler dataMessageHandler;
+
+  [[[self.mockRmqManager stub] andDo:^(NSInvocation *invocation) {
+    dataMessageHandler([FIRMessagingGetRmq2Id(firstMessageStanza) longLongValue],
+                       firstMessageStanza);
+  }]
+      scanWithRmqMessageHandler:[OCMArg isNil]
+             dataMessageHandler:[OCMArg checkWithBlock:^BOOL(id obj) {
+               dataMessageHandler = obj;
+               return YES;
+             }]];
+
+  // expect both 1 and 2 messages to be sent once we regain connection
+  __block BOOL firstMessageSent = NO;
+  __block BOOL secondMessageSent = NO;
+  XCTestExpectation *didSendAllMessages =
+      [self expectationWithDescription:@"Did send all messages"];
+  OCMExpect([self.mockClient sendMessage:[OCMArg checkWithBlock:^BOOL(id obj) {
+    // [didSendAllMessages fulfill];
+    if ([obj isKindOfClass:[GtalkDataMessageStanza class]]) {
+      GtalkDataMessageStanza *message = (GtalkDataMessageStanza *)obj;
+      if ([@"1" isEqualToString:message.id_p]) {
+        firstMessageSent = YES;
+      } else if ([@"2" isEqualToString:message.id_p]) {
+        secondMessageSent = YES;
+      }
+      if (firstMessageSent && secondMessageSent) {
+        [didSendAllMessages fulfill];
+      }
+      return firstMessageSent || secondMessageSent;
+    }
+    return NO;
+  }]]);
+
+  // send the second message after some delay
+  [NSThread sleepForTimeInterval:2.0];
+
+  isConnectionActive = YES;
+  // simulate active connection
+  NSString *newMessageID = @"2";
+  NSMutableDictionary *newMessage = [self upstreamMessageWithID:newMessageID
+                                                            ttl:0
+                                                          delay:0];
+  // send another message to resend not sent messages
+  [self.dataMessageManager sendDataMessageStanza:newMessage];
+
+  [self waitForExpectationsWithTimeout:5.0 handler:^(NSError *error) {
+    XCTAssertNil(error);
+    OCMVerifyAll(self.mockClient);
+    OCMVerifyAll(self.mockReceiver);
+  }];
+}
+
+- (void)testSendDelayedMessage_shouldNotSend {
+  // should not send a delayed message even with an active connection
+  // simulate active connection
+  [[[self.mockClient stub] andReturnValue:OCMOCK_VALUE(YES)] isConnectionActive];
+  [[self.mockClient reject] sendMessage:[OCMArg any]];
+
+  [[self.mockReceiver reject] didSendDataMessageWithID:[OCMArg any]];
+
+  // delayed message
+  NSString *messageID = @"1";
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  NSMutableDictionary *message = [self upstreamMessageWithID:messageID ttl:0 delay:1];
+  [self.dataMessageManager sendDataMessageStanza:message];
+
+  OCMVerifyAll(self.mockClient);
+  OCMVerifyAll(self.mockReceiver);
+}
+
+- (void)testProcessPacket_withValidPacket {
+  GtalkDataMessageStanza *message = [self validDataMessagePacket];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:message];
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingFromKey], message.from);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingCollapseKey], message.token);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingMessageIDKey], kMessagePersistentID);
+  XCTAssertEqualObjects(parsedMessage[kAppDataItemKey], kAppDataItemValue);
+  XCTAssertEqual(4, parsedMessage.count);
+}
+
+- (void)testProcessPacket_withOnlyFrom {
+  GtalkDataMessageStanza *message = [self validDataMessageWithOnlyFrom];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:message];
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingFromKey], message.from);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingMessageIDKey], kMessagePersistentID);
+  XCTAssertEqual(2, parsedMessage.count);
+}
+
+- (void)testProcessPacket_withInvalidPacket {
+  GtalkDataMessageStanza *message = [self invalidDataMessageUsingReservedKeyword];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:message];
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingFromKey], message.from);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingMessageIDKey], kMessagePersistentID);
+  XCTAssertEqual(2, parsedMessage.count);
+}
+
+/**
+ *  Test parsing a duplex message.
+ */
+- (void)testProcessPacket_withDuplexMessage {
+  GtalkDataMessageStanza *stanza = [self validDuplexmessage];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:stanza];
+  XCTAssertEqual(5, parsedMessage.count);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingFromKey], stanza.from);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingCollapseKey], stanza.token);
+  XCTAssertEqualObjects(parsedMessage[kFIRMessagingMessageIDKey], kMessagePersistentID);
+  XCTAssertEqualObjects(parsedMessage[kAppDataItemKey], kAppDataItemValue);
+  XCTAssertTrue([parsedMessage[kFIRMessagingMessageSyncViaMCSKey] boolValue]);
+}
+
+- (void)testReceivingParsedMessage {
+  NSDictionary *message = @{ @"hello" : @"world" };
+  OCMStub([self.mockReceiver didReceiveMessage:[OCMArg isEqual:message] withIdentifier:[OCMArg any]]);
+  [self.dataMessageManager didReceiveParsedMessage:message];
+  OCMVerify([self.mockReceiver didReceiveMessage:message withIdentifier:[OCMArg any]]);
+}
+
+/**
+ *  Test receiving a new duplex message notifies the receiver callback.
+ */
+- (void)testReceivingNewDuplexMessage {
+  GtalkDataMessageStanza *message = [self validDuplexmessage];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:message];
+  [[[self.mockSyncMessageManager stub] andReturnValue:@(NO)]
+      didReceiveMCSSyncMessage:parsedMessage];
+  OCMStub([self.mockReceiver didReceiveMessage:[OCMArg isEqual:message] withIdentifier:[OCMArg any]]);
+  [self.dataMessageManager didReceiveParsedMessage:parsedMessage];
+  OCMVerify([self.mockReceiver didReceiveMessage:[OCMArg any] withIdentifier:[OCMArg any]]);
+}
+
+/**
+ *  Test receiving a duplicated duplex message does not notify the receiver callback.
+ */
+- (void)testReceivingDuplicateDuplexMessage {
+  GtalkDataMessageStanza *message = [self validDuplexmessage];
+  NSDictionary *parsedMessage = [self.dataMessageManager processPacket:message];
+  [[[self.mockSyncMessageManager stub] andReturnValue:@(YES)]
+      didReceiveMCSSyncMessage:parsedMessage];
+  [[self.mockReceiver reject] didReceiveMessage:[OCMArg any] withIdentifier:[OCMArg any]];
+  [self.dataMessageManager didReceiveParsedMessage:parsedMessage];
+}
+
+/**
+ *  In this test we simulate a real RMQ manager and send messages simulating no
+ *  active connection. Then we simulate a new connection being established and
+ *  the client receives a Streaming ACK which should result in resending RMQ messages.
+ */
+- (void)testResendSavedMessages {
+  static BOOL isClientConnected = NO;
+  [[[self.mockClient stub] andDo:^(NSInvocation *invocation) {
+    [invocation setReturnValue:&isClientConnected];
+  }] isConnectionActive];
+
+  // Set a fake, valid bundle identifier
+  [[[self.mockDataMessageManager stub] andReturn:@"gcm-dmm-test"] categoryForUpstreamMessages];
+
+  [FIRMessagingRmqManager removeDatabaseWithName:kRmqDatabaseName];
+  FIRMessagingRmqManager *newRmqManager =
+      [[FIRMessagingRmqManager alloc] initWithDatabaseName:kRmqDatabaseName];
+  [newRmqManager loadRmqId];
+  // have a real RMQ store
+  [self.dataMessageManager setRmq2Manager:newRmqManager];
+
+  [self addFakeFIRMessagingRegistrationToken];
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+
+  // send a couple of message with no connection should be saved to RMQ
+  [self.dataMessageManager sendDataMessageStanza:
+      [self upstreamMessageWithID:@"1" ttl:20000 delay:0]];
+  [self.dataMessageManager sendDataMessageStanza:
+      [self upstreamMessageWithID:@"2" ttl:20000 delay:0]];
+
+  [NSThread sleepForTimeInterval:1.0];
+  isClientConnected = YES;
+  // after the usual version, login assertion we would receive a SelectiveAck
+  // assuming we we weren't able to send any messages we won't delete anything
+  // from the RMQ but try to resend whatever is there
+  __block int didRecieveMessages = 0;
+  id mockConnection = OCMClassMock([FIRMessagingConnection class]);
+
+  BOOL (^resendMessageBlock)(id obj) = ^BOOL(id obj) {
+    if ([obj isKindOfClass:[GtalkDataMessageStanza class]]) {
+      GtalkDataMessageStanza *message = (GtalkDataMessageStanza *)obj;
+      NSLog(@"hello resending %@, %d", message.id_p, didRecieveMessages);
+      if ([@"1" isEqualToString:message.id_p]) {
+        didRecieveMessages |= 1; // right most bit for 1st message
+        return YES;
+      } else if ([@"2" isEqualToString:message.id_p]) {
+        didRecieveMessages |= (1<<1); // second from RMB for 2nd message
+        return YES;
+      }
+    }
+    return NO;
+  };
+  [[[mockConnection stub] andDo:^(NSInvocation *invocation) {
+    // pass
+  }] sendProto:[OCMArg checkWithBlock:resendMessageBlock]];
+
+  [self.dataMessageManager resendMessagesWithConnection:mockConnection];
+
+  // should send both messages
+  XCTAssert(didRecieveMessages == 3);
+  OCMVerifyAll(mockConnection);
+}
+
+- (void)testResendingExpiredMessagesFails {
+  // TODO: Test that expired messages should not be sent on resend
+  static BOOL isClientConnected = NO;
+  [[[self.mockClient stub] andDo:^(NSInvocation *invocation) {
+    [invocation setReturnValue:&isClientConnected];
+  }] isConnectionActive];
+
+  // Set a fake, valid bundle identifier
+  [[[self.mockDataMessageManager stub] andReturn:@"gcm-dmm-test"] categoryForUpstreamMessages];
+
+  [FIRMessagingRmqManager removeDatabaseWithName:kRmqDatabaseName];
+  FIRMessagingRmqManager *newRmqManager =
+      [[FIRMessagingRmqManager alloc] initWithDatabaseName:kRmqDatabaseName];
+  [newRmqManager loadRmqId];
+  // have a real RMQ store
+  [self.dataMessageManager setRmq2Manager:newRmqManager];
+
+  [self.dataMessageManager setDeviceAuthID:@"auth-id" secretToken:@"secret-token"];
+  // send a message that expires in 1 sec
+  [self.dataMessageManager sendDataMessageStanza:
+      [self upstreamMessageWithID:@"1" ttl:1 delay:0]];
+
+  // wait for 2 seconds (let the above message expire)
+  [NSThread sleepForTimeInterval:2.0];
+  isClientConnected = YES;
+
+  id mockConnection = OCMClassMock([FIRMessagingConnection class]);
+
+  [[mockConnection reject] sendProto:[OCMArg any]];
+  [self.dataMessageManager resendMessagesWithConnection:mockConnection];
+
+  // rmq should not have any pending messages
+  [newRmqManager scanWithRmqMessageHandler:^(int64_t rmqId, int8_t tag, NSData *data) {
+    XCTFail(@"RMQ should not have any message");
+  }
+                        dataMessageHandler:nil];
+}
+
+#pragma mark - Private
+
+- (void)addFakeFIRMessagingRegistrationToken {
+  // [[FIRMessagingDefaultsManager sharedInstance] saveAppIDToken:kFIRMessagingAppIDToken];
+}
+
+#pragma mark - Create Packet
+
+- (GtalkDataMessageStanza *)validDataMessagePacket {
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  message.from = kMessageFrom;
+  message.token = kCollapseKey;
+  message.persistentId = kMessagePersistentID;
+  GtalkAppData *item = [[GtalkAppData alloc] init];
+  item.key = kAppDataItemKey;
+  item.value = kAppDataItemValue;
+  message.appDataArray = [NSMutableArray arrayWithObject:item];
+  return message;
+}
+
+- (GtalkDataMessageStanza *)validDataMessageWithOnlyFrom {
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  message.from = kMessageFrom;
+  message.persistentId = kMessagePersistentID;
+  return message;
+}
+
+- (GtalkDataMessageStanza *)invalidDataMessageUsingReservedKeyword {
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  message.from = kMessageFrom;
+  message.persistentId = kMessagePersistentID;
+  GtalkAppData *item = [[GtalkAppData alloc] init];
+  item.key = kAppDataItemInvalidKey;
+  item.value = kAppDataItemValue;
+  message.appDataArray = [NSMutableArray arrayWithObject:item];
+  return message;
+}
+
+- (GtalkDataMessageStanza *)validDataMessageForFIRMessaging {
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  message.from = kMessageFrom;
+  message.token = @"com.google.gcm";
+  return message;
+}
+
+- (GtalkDataMessageStanza *)validDuplexmessage {
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  message.from = kMessageFrom;
+  message.token = kCollapseKey;
+  message.persistentId = kMessagePersistentID;
+  GtalkAppData *item = [[GtalkAppData alloc] init];
+  item.key = kAppDataItemKey;
+  item.value = kAppDataItemValue;
+  GtalkAppData *duplexItem = [[GtalkAppData alloc] init];
+  duplexItem.key = @"gcm.duplex";
+  duplexItem.value = @"1";
+  message.appDataArray = [NSMutableArray arrayWithObjects:item, duplexItem, nil];
+  return message;
+}
+
+#pragma mark - Create Message
+
+- (NSMutableDictionary *)standardFIRMessagingMessageWithMessageID:(NSString *)messageID {
+  NSDictionary *message = @{ kAppDataItemKey : kAppDataItemValue };
+  return [FIRMessaging createFIRMessagingMessageWithMessage:message
+                                                         to:kMessageTo
+                                                     withID:messageID
+                                                 timeToLive:-1
+                                                      delay:0];
+}
+
+- (NSMutableDictionary *)upstreamMessageWithID:(NSString *)messageID
+                                           ttl:(int64_t)ttl
+                                         delay:(int)delay {
+  NSDictionary *message = @{ kAppDataItemInvalidKey : kAppDataItemValue };
+  return [FIRMessaging createFIRMessagingMessageWithMessage:message
+                                                         to:kMessageTo
+                                                     withID:messageID
+                                                 timeToLive:ttl
+                                                      delay:delay];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingFakeConnection.h b/Example/Messaging/Tests/FIRMessagingFakeConnection.h
new file mode 100644
index 0000000..7fe52bf
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingFakeConnection.h
@@ -0,0 +1,63 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FIRMessagingConnection.h"
+
+/**
+ * A bunch of different fake connections are used to simulate various connection behaviours.
+ * A fake connection that successfully conects to remote host.
+ */
+// TODO: Split FIRMessagingConnection to make it more testable.
+@interface FIRMessagingFakeConnection : FIRMessagingConnection
+
+@property(nonatomic, readwrite, assign) BOOL shouldFakeSuccessLogin;
+
+// timeout caused by heartbeat failure (defaults to 0.5s)
+@property(nonatomic, readwrite, assign) NSTimeInterval fakeConnectionTimeout;
+
+/**
+ * Should stub the socket disconnect to not fail when called
+ */
+- (void)mockSocketDisconnect;
+
+/**
+ * Calls disconnect on the socket(which should theoretically be mocked by the above method) and
+ * let the socket delegate know that it has been disconnected.
+ */
+- (void)disconnectNow;
+
+/**
+ * The fake host to connect to.
+ */
++ (NSString *)fakeHost;
+
+/**
+ * The fake port used to connect.
+ */
++ (int)fakePort;
+
+@end
+
+/**
+ * A fake connection that simulates failure a certain number of times before success.
+ */
+// TODO: Coalesce this with the FIRMessagingFakeConnection itself.
+@interface FIRMessagingFakeFailConnection : FIRMessagingFakeConnection
+
+@property(nonatomic, readwrite, assign) int failCount;
+@property(nonatomic, readwrite, assign) int signInRequests;
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingFakeConnection.m b/Example/Messaging/Tests/FIRMessagingFakeConnection.m
new file mode 100644
index 0000000..de2e0bb
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingFakeConnection.m
@@ -0,0 +1,150 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FIRMessagingFakeConnection.h"
+
+#import <OCMock/OCMock.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+
+#import "FIRMessagingSecureSocket.h"
+#import "FIRMessagingUtilities.h"
+
+static NSString *const kHost = @"localhost";
+static const int kPort = 6234;
+
+@interface FIRMessagingSecureSocket ()
+
+@property(nonatomic, readwrite, assign) FIRMessagingSecureSocketState state;
+
+@end
+
+@interface FIRMessagingConnection ()
+
+@property(nonatomic, readwrite, strong) FIRMessagingSecureSocket *socket;
+
+- (void)setupConnectionSocket;
+- (void)connectToSocket:(FIRMessagingSecureSocket *)socket;
+- (NSTimeInterval)connectionTimeoutInterval;
+- (void)sendHeartbeatPing;
+- (void)secureSocket:(FIRMessagingSecureSocket *)socket
+      didReceiveData:(NSData *)data
+             withTag:(int8_t)tag;
+
+@end
+
+@implementation FIRMessagingFakeConnection
+
+- (void)signIn {
+
+  // use this if you don't really want to mock/stub the login behaviour. In case
+  // you want to stub the login behavoiur you should do these things manually in
+  // your test and add custom logic in between as required for your testing.
+  [self setupConnectionSocket];
+
+  id socketMock = OCMPartialMock(self.socket);
+  self.socket = socketMock;
+  [[[socketMock stub]
+      andDo:^(NSInvocation *invocation) {
+        if (self.shouldFakeSuccessLogin) {
+          [self willFakeSuccessfulLoginToFCM];
+        }
+        self.socket.state = kFIRMessagingSecureSocketOpen;
+        [self.socket.delegate secureSocketDidConnect:self.socket];
+      }]
+      connectToHost:kHost
+               port:kPort
+          onRunLoop:[OCMArg any]];
+
+  [self connectToSocket:socketMock];
+}
+
+- (NSTimeInterval)connectionTimeoutInterval {
+  if (self.fakeConnectionTimeout) {
+    return self.fakeConnectionTimeout;
+  } else {
+    return 0.5;  // 0.5s
+  }
+}
+
+- (void)mockSocketDisconnect {
+  id mockSocket = self.socket;
+  [[[mockSocket stub] andDo:^(NSInvocation *invocation) {
+    self.socket.state = kFIRMessagingSecureSocketClosed;
+  }] disconnect];
+}
+
+- (void)disconnectNow {
+  [self.socket disconnect];
+  [self.socket.delegate didDisconnectWithSecureSocket:self.socket];
+}
+
++ (NSString *)fakeHost {
+  return @"localhost";
+}
+
++ (int)fakePort {
+  return 6234;
+}
+
+- (void)willFakeSuccessfulLoginToFCM {
+  id mockSocket = self.socket;
+  [[[mockSocket stub]
+      andDo:^(NSInvocation *invocation) {
+        // mock successful login
+
+        GtalkLoginResponse *response = [[GtalkLoginResponse alloc] init];
+        [response setId_p:@""];
+        [self secureSocket:self.socket
+            didReceiveData:[response data]
+                   withTag:kFIRMessagingProtoTagLoginResponse];
+      }]
+      sendData:[OCMArg any]
+       withTag:kFIRMessagingProtoTagLoginRequest
+         rmqId:[OCMArg isNil]];
+}
+
+@end
+
+@implementation FIRMessagingFakeFailConnection
+
+- (void)signIn {
+  self.signInRequests++;
+  [self setupConnectionSocket];
+  id mockSocket = OCMPartialMock(self.socket);
+  self.socket = mockSocket;
+  [[[mockSocket stub]
+      andDo:^(NSInvocation *invocation) {
+        [self mockSocketDisconnect];
+        if (self.signInRequests <= self.failCount) {
+          // do nothing -- should timeout
+        } else {
+          // since we will always fail once we would disconnect the socket before
+          // we ever try again thus mock the disconnect to change the state and
+          // prevent any assertions
+          [self willFakeSuccessfulLoginToFCM];
+          self.socket.state = kFIRMessagingSecureSocketOpen;
+          [self.socket.delegate secureSocketDidConnect:self.socket];
+        }
+      }]
+      connectToHost:kHost
+               port:kPort
+          onRunLoop:[OCMArg any]];
+
+  [self connectToSocket:mockSocket];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingFakeSocket.h b/Example/Messaging/Tests/FIRMessagingFakeSocket.h
new file mode 100644
index 0000000..ecba9dc
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingFakeSocket.h
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FIRMessagingSecureSocket.h"
+
+@interface FIRMessagingFakeSocket : FIRMessagingSecureSocket
+
+/**
+ * Initialize socket with a given buffer size. Designated Initializer.
+ *
+ *  @param bufferSize  The buffer size used to connect the input and the output stream. Note
+ *                     when we write data to the output stream it's read in terms of this buffer
+ *                     size. So for tests using `FIRMessagingFakeSocket` you should use an appropriate
+ *                     buffer size in terms of what you are writing to the buffer and what should
+ *                     be read. Since there is no "flush" operation in NSStream we would have to
+ *                     live with this.
+ *
+ *  @see {FIRMessagingSecureSocketTest} for example usage.
+ *  @return A fake secure socket.
+ */
+- (instancetype)initWithBufferSize:(uint8_t)bufferSize;
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingFakeSocket.m b/Example/Messaging/Tests/FIRMessagingFakeSocket.m
new file mode 100644
index 0000000..2b0b477
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingFakeSocket.m
@@ -0,0 +1,89 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FIRMessagingFakeSocket.h"
+
+#import "FIRMessagingConstants.h"
+#import "FIRMessagingDefines.h"
+#import <GoogleToolboxForMac/GTMDefines.h>
+
+@interface FIRMessagingSecureSocket() <NSStreamDelegate>
+
+@property(nonatomic, readwrite, assign) FIRMessagingSecureSocketState state;
+@property(nonatomic, readwrite, strong) NSInputStream *inStream;
+@property(nonatomic, readwrite, strong) NSOutputStream *outStream;
+
+@property(nonatomic, readwrite, assign) BOOL isInStreamOpen;
+@property(nonatomic, readwrite, assign) BOOL isOutStreamOpen;
+
+@property(nonatomic, readwrite, strong) NSRunLoop *runLoop;
+
+@end
+
+@interface FIRMessagingFakeSocket ()
+
+@property(nonatomic, readwrite, assign) int8_t bufferSize;
+
+@end
+
+@implementation FIRMessagingFakeSocket
+
+- (instancetype)initWithBufferSize:(uint8_t)bufferSize {
+  self = [super init];
+  if (self) {
+    _bufferSize = bufferSize;
+  }
+  return self;
+}
+
+- (void)connectToHost:(NSString *)host
+                 port:(NSUInteger)port
+            onRunLoop:(NSRunLoop *)runLoop {
+  self.state = kFIRMessagingSecureSocketOpening;
+  self.runLoop = runLoop;
+
+  CFReadStreamRef inputStreamRef = nil;
+  CFWriteStreamRef outputStreamRef = nil;
+
+  CFStreamCreateBoundPair(NULL,
+                          &inputStreamRef,
+                          &outputStreamRef,
+                          self.bufferSize);
+
+  self.inStream = CFBridgingRelease(inputStreamRef);
+  self.outStream = CFBridgingRelease(outputStreamRef);
+  if (!self.inStream || !self.outStream) {
+    FIRMessaging_FAIL(@"cannot create a fake socket");
+    return;
+  }
+
+  self.isInStreamOpen = NO;
+  self.isOutStreamOpen = NO;
+
+  [self openStream:self.outStream];
+  [self openStream:self.inStream];
+}
+
+- (void)openStream:(NSStream *)stream {
+  _GTMDevAssert(stream, @"Cannot open nil stream");
+  if (stream) {
+    stream.delegate = self;
+    [stream scheduleInRunLoop:self.runLoop forMode:NSDefaultRunLoopMode];
+    [stream open];
+  }
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingLinkHandlingTest.m b/Example/Messaging/Tests/FIRMessagingLinkHandlingTest.m
new file mode 100644
index 0000000..94ce530
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingLinkHandlingTest.m
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessaging.h"
+#import "FIRMessagingConfig.h"
+#import "FIRMessagingConstants.h"
+#import "FIRMessagingTestNotificationUtilities.h"
+
+@interface FIRMessaging ()
+
+- (instancetype)initWithConfig:(FIRMessagingConfig *)config;
+- (NSURL *)linkURLFromMessage:(NSDictionary *)message;
+
+@end
+
+@interface FIRMessagingLinkHandlingTest : XCTestCase
+
+@property(nonatomic, readonly, strong) FIRMessaging *messaging;
+
+@end
+
+@implementation FIRMessagingLinkHandlingTest
+
+- (void)setUp {
+  [super setUp];
+
+  FIRMessagingConfig *config = [FIRMessagingConfig defaultConfig];
+  _messaging = [[FIRMessaging alloc] initWithConfig:config];
+}
+
+- (void)tearDown {
+  _messaging = nil;
+  [super tearDown];
+}
+
+#pragma mark - Link Handling Testing
+
+- (void)testNonExistentLinkInMessage {
+  NSMutableDictionary *notification =
+      [FIRMessagingTestNotificationUtilities createBasicNotificationWithUniqueMessageID];
+  NSURL *url = [self.messaging linkURLFromMessage:notification];
+  XCTAssertNil(url);
+}
+
+- (void)testEmptyLinkInMessage {
+  NSMutableDictionary *notification =
+      [FIRMessagingTestNotificationUtilities createBasicNotificationWithUniqueMessageID];
+  notification[kFIRMessagingMessageLinkKey] = @"";
+  NSURL *url = [self.messaging linkURLFromMessage:notification];
+  XCTAssertNil(url);
+}
+
+- (void)testNonStringLinkInMessage {
+  NSMutableDictionary *notification =
+      [FIRMessagingTestNotificationUtilities createBasicNotificationWithUniqueMessageID];
+  notification[kFIRMessagingMessageLinkKey] = @(5);
+  NSURL *url = [self.messaging linkURLFromMessage:notification];
+  XCTAssertNil(url);
+}
+
+- (void)testInvalidURLStringLinkInMessage {
+  NSMutableDictionary *notification =
+      [FIRMessagingTestNotificationUtilities createBasicNotificationWithUniqueMessageID];
+  notification[kFIRMessagingMessageLinkKey] = @"This is not a valid url string";
+  NSURL *url = [self.messaging linkURLFromMessage:notification];
+  XCTAssertNil(url);
+}
+
+- (void)testValidURLStringLinkInMessage {
+  NSMutableDictionary *notification =
+      [FIRMessagingTestNotificationUtilities createBasicNotificationWithUniqueMessageID];
+  notification[kFIRMessagingMessageLinkKey] = @"https://www.google.com/";
+  NSURL *url = [self.messaging linkURLFromMessage:notification];
+  XCTAssertTrue([url.absoluteString isEqualToString:@"https://www.google.com/"]);
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingPendingTopicsListTest.m b/Example/Messaging/Tests/FIRMessagingPendingTopicsListTest.m
new file mode 100644
index 0000000..2033cb4
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingPendingTopicsListTest.m
@@ -0,0 +1,263 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <OCMock/OCMock.h>
+#import <XCTest/XCTest.h>
+
+#import "FIRMessagingDefines.h"
+#import "FIRMessagingPendingTopicsList.h"
+#import "FIRMessagingTopicsCommon.h"
+
+typedef void (^MockDelegateSubscriptionHandler)(NSString *topic,
+                                                FIRMessagingTopicAction action,
+                                                FIRMessagingTopicOperationCompletion completion);
+
+/**
+ * This object lets us provide a stub delegate where we can customize the behavior by providing
+ * blocks. We need to use this instead of stubbing a OCMockProtocol because our delegate methods
+ * take primitive values (e.g. action), which is not easy to use from OCMock
+ * @see http://stackoverflow.com/a/6332023
+ */
+@interface MockPendingTopicsListDelegate: NSObject <FIRMessagingPendingTopicsListDelegate>
+
+@property(nonatomic, assign) BOOL isReady;
+@property(nonatomic, copy) MockDelegateSubscriptionHandler subscriptionHandler;
+@property(nonatomic, copy) void(^updateHandler)();
+
+@end
+
+@implementation MockPendingTopicsListDelegate
+
+- (BOOL)pendingTopicsListCanRequestTopicUpdates:(FIRMessagingPendingTopicsList *)list {
+  return self.isReady;
+}
+
+- (void)pendingTopicsList:(FIRMessagingPendingTopicsList *)list
+  requestedUpdateForTopic:(NSString *)topic
+                   action:(FIRMessagingTopicAction)action
+               completion:(FIRMessagingTopicOperationCompletion)completion {
+  if (self.subscriptionHandler) {
+    self.subscriptionHandler(topic, action, completion);
+  }
+}
+
+- (void)pendingTopicsListDidUpdate:(FIRMessagingPendingTopicsList *)list {
+  if (self.updateHandler) {
+    self.updateHandler();
+  }
+}
+
+@end
+
+@interface FIRMessagingPendingTopicsListTest : XCTestCase
+
+/// Using this delegate lets us prevent any topic operations from start, making it easy to measure
+/// our batches
+@property(nonatomic, strong) MockPendingTopicsListDelegate *notReadyDelegate;
+/// Using this delegate will always begin topic operations (which will never return by default).
+/// Useful for overriding with block-methods to handle update requests
+@property(nonatomic, strong) MockPendingTopicsListDelegate *alwaysReadyDelegate;
+
+@end
+
+@implementation FIRMessagingPendingTopicsListTest
+
+- (void)setUp {
+  [super setUp];
+  self.notReadyDelegate = [[MockPendingTopicsListDelegate alloc] init];
+  self.notReadyDelegate.isReady = NO;
+
+  self.alwaysReadyDelegate = [[MockPendingTopicsListDelegate alloc] init];
+  self.alwaysReadyDelegate.isReady = YES;
+}
+
+- (void)tearDown {
+  self.notReadyDelegate = nil;
+  self.alwaysReadyDelegate = nil;
+  [super tearDown];
+}
+
+- (void)testAddSingleTopic {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.notReadyDelegate;
+
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  XCTAssertEqual(pendingTopics.numberOfBatches, 1);
+}
+
+- (void)testAddSameTopicAndActionMultipleTimes {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.notReadyDelegate;
+
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  XCTAssertEqual(pendingTopics.numberOfBatches, 1);
+}
+
+- (void)testAddMultiplePendingTopicsWithSameAction {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.notReadyDelegate;
+
+  for (NSInteger i = 0; i < 10; i++) {
+    NSString *topic = [NSString stringWithFormat:@"/topics/%ld", i];
+    [pendingTopics addOperationForTopic:topic
+                             withAction:FIRMessagingTopicActionSubscribe
+                             completion:nil];
+  }
+  XCTAssertEqual(pendingTopics.numberOfBatches, 1);
+}
+
+- (void)testAddTopicsWithDifferentActions {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.notReadyDelegate;
+
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/1"
+                           withAction:FIRMessagingTopicActionUnsubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/2"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  XCTAssertEqual(pendingTopics.numberOfBatches, 3);
+}
+
+- (void)testBatchSizeReductionAfterSuccessfulTopicUpdate {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.alwaysReadyDelegate;
+
+  XCTestExpectation *batchSizeReductionExpectation =
+      [self expectationWithDescription:@"Batch size was reduced after topic suscription"];
+
+  FIRMessaging_WEAKIFY(self)
+  self.alwaysReadyDelegate.subscriptionHandler =
+      ^(NSString *topic,
+        FIRMessagingTopicAction action,
+        FIRMessagingTopicOperationCompletion completion) {
+    // Simulate that the handler is generally called asynchronously
+    dispatch_async(dispatch_get_main_queue(), ^{
+      FIRMessaging_STRONGIFY(self)
+      if (action == FIRMessagingTopicActionUnsubscribe) {
+        XCTAssertEqual(pendingTopics.numberOfBatches, 1);
+        [batchSizeReductionExpectation fulfill];
+      }
+      completion(FIRMessagingTopicOperationResultSucceeded, nil);
+    });
+  };
+
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/1"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/2"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/1"
+                           withAction:FIRMessagingTopicActionUnsubscribe
+                           completion:nil];
+
+  [self waitForExpectationsWithTimeout:5.0 handler:nil];
+}
+
+- (void)testCompletionOfTopicUpdatesInSameThread {
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.alwaysReadyDelegate;
+
+  XCTestExpectation *allOperationsSucceededed =
+      [self expectationWithDescription:@"All queued operations succeeded"];
+
+  self.alwaysReadyDelegate.subscriptionHandler =
+      ^(NSString *topic,
+        FIRMessagingTopicAction action,
+        FIRMessagingTopicOperationCompletion completion) {
+    // Typically, our callbacks happen asynchronously, but to ensure resilience,
+    // call back the operation on the same thread it was called in.
+    completion(FIRMessagingTopicOperationResultSucceeded, nil);
+  };
+
+  self.alwaysReadyDelegate.updateHandler = ^{
+    if (pendingTopics.numberOfBatches == 0) {
+      [allOperationsSucceededed fulfill];
+    }
+  };
+
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/1"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  [pendingTopics addOperationForTopic:@"/topics/2"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+
+  [self waitForExpectationsWithTimeout:5.0 handler:nil];
+}
+
+- (void)testAddingTopicToCurrentBatchWhileCurrentBatchTopicsInFlight {
+
+  FIRMessagingPendingTopicsList *pendingTopics = [[FIRMessagingPendingTopicsList alloc] init];
+  pendingTopics.delegate = self.alwaysReadyDelegate;
+
+  NSString *stragglerTopic = @"/topics/straggler";
+  XCTestExpectation *stragglerTopicWasAddedToInFlightOperations =
+  [self expectationWithDescription:@"The topic was added to in-flight operations"];
+
+  self.alwaysReadyDelegate.subscriptionHandler =
+      ^(NSString *topic,
+        FIRMessagingTopicAction action,
+        FIRMessagingTopicOperationCompletion completion) {
+    if ([topic isEqualToString:stragglerTopic]) {
+      [stragglerTopicWasAddedToInFlightOperations fulfill];
+    }
+    // Add a 0.5 second delay to the completion, to give time to add a straggler before the batch
+    // is completed
+    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.5 * NSEC_PER_SEC)),
+                   dispatch_get_main_queue(),
+                   ^{
+      completion(FIRMessagingTopicOperationResultSucceeded, nil);
+    });
+  };
+
+  // This is a normal topic, which should start fairly soon, but take a while to complete
+  [pendingTopics addOperationForTopic:@"/topics/0"
+                           withAction:FIRMessagingTopicActionSubscribe
+                           completion:nil];
+  // While waiting for the first topic to complete, we add another topic after a slight delay
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.25 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(),
+                 ^{
+    [pendingTopics addOperationForTopic:stragglerTopic
+                             withAction:FIRMessagingTopicActionSubscribe
+                             completion:nil];
+  });
+
+  [self waitForExpectationsWithTimeout:5.0 handler:nil];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingPubSubTest.m b/Example/Messaging/Tests/FIRMessagingPubSubTest.m
new file mode 100644
index 0000000..2981b54
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingPubSubTest.m
@@ -0,0 +1,81 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import "FIRMessagingPubSub.h"
+
+@interface FIRMessagingPubSubTest : XCTestCase
+@end
+
+@implementation FIRMessagingPubSubTest
+
+static NSString *const kTopicName = @"topic-Name";
+
+#pragma mark - topicMatchForSender tests
+
+/// Tests that an empty topic name is an invalid topic.
+- (void)testTopicMatchForEmptyTopicPrefix {
+  XCTAssertFalse([FIRMessagingPubSub isValidTopicWithPrefix:@""]);
+}
+
+/// Tests that a topic with an invalid prefix is not a valid topic name.
+- (void)testTopicMatchWithInvalidTopicPrefix {
+  XCTAssertFalse([FIRMessagingPubSub isValidTopicWithPrefix:@"/topics+abcdef/"]);
+}
+
+/// Tests that a topic with a valid prefix but invalid name is not a valid topic name.
+- (void)testTopicMatchWithValidTopicPrefixButInvalidName {
+  XCTAssertFalse([FIRMessagingPubSub isValidTopicWithPrefix:@"/topics/aaaaaa/topics/lala"]);
+}
+
+/// Tests that multiple backslashes in topics is an invalid topic name.
+- (void)testTopicMatchForInvalidTopicPrefix_multipleBackslash {
+  XCTAssertFalse([FIRMessagingPubSub isValidTopicWithPrefix:@"/topics//abc"]);
+}
+
+/// Tests a topic name with a valid prefix and name.
+- (void)testTopicMatchForValidTopicSender {
+  NSString *topic = [NSString stringWithFormat:@"/topics/%@", kTopicName];
+  XCTAssertTrue([FIRMessagingPubSub isValidTopicWithPrefix:topic]);
+}
+
+/// Tests topic prefix for topics with no prefix.
+- (void)testTopicHasNoTopicPrefix {
+  XCTAssertFalse([FIRMessagingPubSub hasTopicsPrefix:@""]);
+}
+
+/// Tests topic prefix for valid prefix.
+- (void)testTopicHasValidToicsPrefix {
+  XCTAssertTrue([FIRMessagingPubSub hasTopicsPrefix:@"/topics/"]);
+}
+
+/// Tests topic prefix wih no prefix.
+- (void)testAddTopicPrefix_withNoPrefix {
+  NSString *topic = [FIRMessagingPubSub addPrefixToTopic:@""];
+  XCTAssertTrue([FIRMessagingPubSub hasTopicsPrefix:topic]);
+  XCTAssertFalse([FIRMessagingPubSub isValidTopicWithPrefix:topic]);
+}
+
+/// Tests adding the "/topics/" prefix for topic name which already has a prefix.
+- (void)testAddTopicPrefix_withPrefix {
+  NSString *topic = [NSString stringWithFormat:@"/topics/%@", kTopicName];
+  topic = [FIRMessagingPubSub addPrefixToTopic:topic];
+  XCTAssertTrue([FIRMessagingPubSub hasTopicsPrefix:topic]);
+  XCTAssertTrue([FIRMessagingPubSub isValidTopicWithPrefix:topic]);
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingRegistrarTest.m b/Example/Messaging/Tests/FIRMessagingRegistrarTest.m
new file mode 100644
index 0000000..b32851c
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingRegistrarTest.m
@@ -0,0 +1,134 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import UIKit;
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessagingCheckinService.h"
+#import "FIRMessagingPubSubRegistrar.h"
+#import "FIRMessagingRegistrar.h"
+#import "FIRMessagingUtilities.h"
+#import "NSError+FIRMessaging.h"
+
+static NSString *const kFIRMessagingUserDefaultsSuite =
+    @"FIRMessagingRegistrarTestUserDefaultsSuite";
+
+static NSString *const kDeviceAuthId = @"12345";
+static NSString *const kSecretToken = @"45657809";
+static NSString *const kVersionInfo = @"1.0";
+static NSString *const kTopicToSubscribeTo = @"/topics/xyz/hello-world";
+static NSString *const kFIRMessagingAppIDToken = @"abcdefgh1234lmno";
+static NSString *const kSubscriptionID = @"sample-subscription-id-xyz";
+
+@interface FIRMessagingRegistrar ()
+
+@property(nonatomic, readwrite, strong) FIRMessagingPubSubRegistrar *pubsubRegistrar;
+@property(nonatomic, readwrite, strong) FIRMessagingCheckinService *checkinService;
+
+@end
+
+@interface FIRMessagingRegistrarTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) FIRMessagingRegistrar *registrar;
+@property(nonatomic, readwrite, strong) id mockRegistrar;
+@property(nonatomic, readwrite, strong) id mockCheckin;
+@property(nonatomic, readwrite, strong) id mockPubsubRegistrar;
+
+@end
+
+@implementation FIRMessagingRegistrarTest
+
+- (void)setUp {
+  [super setUp];
+  _registrar = [[FIRMessagingRegistrar alloc] init];
+  _mockRegistrar = OCMPartialMock(_registrar);
+  _mockCheckin = OCMPartialMock(_registrar.checkinService);
+  _registrar.checkinService = _mockCheckin;
+  _registrar.pubsubRegistrar = OCMClassMock([FIRMessagingPubSubRegistrar class]);
+  _mockPubsubRegistrar = _registrar.pubsubRegistrar;
+}
+
+- (void)testUpdateSubscriptionWithValidCheckinData {
+  [self stubCheckinService];
+
+  [self.registrar updateSubscriptionToTopic:kTopicToSubscribeTo
+                                  withToken:kFIRMessagingAppIDToken
+                                    options:nil
+                               shouldDelete:NO
+                                    handler:
+      ^(FIRMessagingTopicOperationResult result, NSError *error) {
+  }];
+
+  OCMVerify([self.mockPubsubRegistrar updateSubscriptionToTopic:[OCMArg isEqual:kTopicToSubscribeTo]
+                                                      withToken:[OCMArg isEqual:kFIRMessagingAppIDToken]
+                                                        options:nil
+                                                   shouldDelete:NO
+                                                        handler:OCMOCK_ANY]);
+}
+
+- (void)testUpdateSubscription {
+  [self stubCheckinService];
+
+  __block FIRMessagingTopicOperationCompletion pubsubCompletion;
+  [[[self.mockPubsubRegistrar stub]
+      andDo:^(NSInvocation *invocation) {
+        pubsubCompletion(FIRMessagingTopicOperationResultSucceeded, nil);
+      }]
+      updateSubscriptionToTopic:kTopicToSubscribeTo
+                      withToken:kFIRMessagingAppIDToken
+                        options:nil
+                   shouldDelete:NO
+                        handler:[OCMArg checkWithBlock:^BOOL(id obj) {
+                          return (pubsubCompletion = obj) != nil;
+                        }]];
+
+  [self.registrar updateSubscriptionToTopic:kTopicToSubscribeTo
+                                  withToken:kFIRMessagingAppIDToken
+                                    options:nil
+                               shouldDelete:NO
+                                    handler:
+      ^(FIRMessagingTopicOperationResult result, NSError *error) {
+    XCTAssertNil(error);
+    XCTAssertEqual(result, FIRMessagingTopicOperationResultSucceeded);
+  }];
+}
+
+- (void)testFailedUpdateSubscriptionWithNoCheckin {
+  // Mock checkin service to always return NO for hasValidCheckinInfo
+  [[[self.mockCheckin stub] andReturnValue:@NO] hasValidCheckinInfo];
+  // This should not create a network request since we don't have checkin info
+  [self.registrar updateSubscriptionToTopic:kTopicToSubscribeTo
+                                  withToken:kFIRMessagingAppIDToken
+                                    options:nil
+                               shouldDelete:NO
+                                    handler:
+      ^(FIRMessagingTopicOperationResult result, NSError *error) {
+    XCTAssertNotNil(error);
+    XCTAssertEqual(result, FIRMessagingTopicOperationResultError);
+  }];
+}
+
+#pragma mark - Private Helpers
+
+- (void)stubCheckinService {
+  [[[self.mockCheckin stub] andReturn:kDeviceAuthId] deviceAuthID];
+  [[[self.mockCheckin stub] andReturn:kSecretToken] secretToken];
+  [[[self.mockCheckin stub] andReturnValue:@YES] hasValidCheckinInfo];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingRemoteNotificationsProxyTest.m b/Example/Messaging/Tests/FIRMessagingRemoteNotificationsProxyTest.m
new file mode 100644
index 0000000..9138c50
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingRemoteNotificationsProxyTest.m
@@ -0,0 +1,279 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#if __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
+@import UserNotifications;
+#endif
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessaging.h"
+#import "FIRMessagingRemoteNotificationsProxy.h"
+
+#pragma mark - Expose Internal Methods for Testing
+// Expose some internal properties and methods here, in order to test
+@interface FIRMessagingRemoteNotificationsProxy ()
+
+@property(readonly, nonatomic) BOOL didSwizzleMethods;
+@property(readonly, nonatomic) BOOL didSwizzleAppDelegateMethods;
+
+@property(readonly, nonatomic) BOOL hasSwizzledUserNotificationDelegate;
+@property(readonly, nonatomic) BOOL isObservingUserNotificationDelegateChanges;
+
+@property(strong, readonly, nonatomic) id userNotificationCenter;
+@property(strong, readonly, nonatomic) id currentUserNotificationCenterDelegate;
+
++ (instancetype)sharedProxy;
+
+- (BOOL)swizzleAppDelegateMethods:(id<UIApplicationDelegate>)appDelegate;
+- (void)listenForDelegateChangesInUserNotificationCenter:(id)notificationCenter;
+- (void)swizzleUserNotificationCenterDelegate:(id)delegate;
+- (void)unswizzleUserNotificationCenterDelegate:(id)delegate;
+
+void FCM_swizzle_appDidReceiveRemoteNotification(id self,
+                                                 SEL _cmd,
+                                                 UIApplication *app,
+                                                 NSDictionary *userInfo);
+void FCM_swizzle_appDidReceiveRemoteNotificationWithHandler(
+    id self, SEL _cmd, UIApplication *app, NSDictionary *userInfo,
+    void (^handler)(UIBackgroundFetchResult));
+void FCM_swizzle_willPresentNotificationWithHandler(
+    id self, SEL _cmd, id center, id notification, void (^handler)(NSUInteger));
+
+@end
+
+#pragma mark - Incomplete App Delegate
+@interface IncompleteAppDelegate : NSObject <UIApplicationDelegate>
+@end
+@implementation IncompleteAppDelegate
+@end
+
+#pragma mark - Fake AppDelegate
+@interface FakeAppDelegate : NSObject <UIApplicationDelegate>
+@property(nonatomic) BOOL remoteNotificationMethodWasCalled;
+@property(nonatomic) BOOL remoteNotificationWithFetchHandlerWasCalled;
+@end
+@implementation FakeAppDelegate
+- (void)application:(UIApplication *)application
+    didReceiveRemoteNotification:(NSDictionary *)userInfo {
+  self.remoteNotificationMethodWasCalled = YES;
+}
+- (void)application:(UIApplication *)application
+    didReceiveRemoteNotification:(NSDictionary *)userInfo
+    fetchCompletionHandler:(void (^)(UIBackgroundFetchResult))completionHandler {
+  self.remoteNotificationWithFetchHandlerWasCalled = YES;
+}
+@end
+
+#pragma mark - Incompete UNUserNotificationCenterDelegate
+@interface IncompleteUserNotificationCenterDelegate : NSObject <UNUserNotificationCenterDelegate>
+@end
+@implementation IncompleteUserNotificationCenterDelegate
+@end
+
+#pragma mark - Fake UNUserNotificationCenterDelegate
+
+#if __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
+@interface FakeUserNotificationCenterDelegate : NSObject <UNUserNotificationCenterDelegate>
+@property(nonatomic) BOOL willPresentWasCalled;
+@end
+@implementation FakeUserNotificationCenterDelegate
+- (void)userNotificationCenter:(UNUserNotificationCenter *)center
+       willPresentNotification:(UNNotification *)notification
+         withCompletionHandler:(void (^)(UNNotificationPresentationOptions options))
+    completionHandler {
+  self.willPresentWasCalled = YES;
+}
+@end
+#endif // __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
+
+#pragma mark - Local, Per-Test Properties
+
+@interface FIRMessagingRemoteNotificationsProxyTest : XCTestCase
+
+@property(nonatomic, strong) FIRMessagingRemoteNotificationsProxy *proxy;
+@property(nonatomic, strong) id mockProxy;
+@property(nonatomic, strong) id mockProxyClass;
+@property(nonatomic, strong) id mockMessagingClass;
+
+@end
+
+@implementation FIRMessagingRemoteNotificationsProxyTest
+
+- (void)setUp {
+  [super setUp];
+  _proxy = [[FIRMessagingRemoteNotificationsProxy alloc] init];
+  _mockProxy = OCMPartialMock(_proxy);
+  _mockProxyClass = OCMClassMock([FIRMessagingRemoteNotificationsProxy class]);
+  // Update +sharedProxy to always return our partial mock of FIRMessagingRemoteNotificationsProxy
+  OCMStub([_mockProxyClass sharedProxy]).andReturn(_mockProxy);
+  // Many of our swizzled methods call [FIRMessaging messaging], but we don't need it,
+  // so just stub it to return nil
+  _mockMessagingClass = OCMClassMock([FIRMessaging class]);
+  OCMStub([_mockMessagingClass messaging]).andReturn(nil);
+}
+
+- (void)tearDown {
+  [_mockMessagingClass stopMocking];
+  _mockMessagingClass = nil;
+
+  [_mockProxyClass stopMocking];
+  _mockProxyClass = nil;
+
+  [_mockProxy stopMocking];
+  _mockProxy = nil;
+
+  _proxy = nil;
+  [super tearDown];
+}
+
+#pragma mark - Method Swizzling Tests
+
+- (void)testSwizzlingNonAppDelegate {
+  id randomObject = @"Random Object that is not an App Delegate";
+  [self.proxy swizzleAppDelegateMethods:randomObject];
+  XCTAssertFalse(self.proxy.didSwizzleAppDelegateMethods);
+}
+
+- (void)testSwizzlingAppDelegate {
+  IncompleteAppDelegate *incompleteAppDelegate = [[IncompleteAppDelegate alloc] init];
+  [self.proxy swizzleAppDelegateMethods:incompleteAppDelegate];
+  XCTAssertTrue(self.proxy.didSwizzleAppDelegateMethods);
+}
+
+- (void)testSwizzledIncompleteAppDelegateRemoteNotificationMethod {
+  IncompleteAppDelegate *incompleteAppDelegate = [[IncompleteAppDelegate alloc] init];
+  [self.mockProxy swizzleAppDelegateMethods:incompleteAppDelegate];
+  SEL selector = @selector(application:didReceiveRemoteNotification:);
+  XCTAssertTrue([incompleteAppDelegate respondsToSelector:selector]);
+  [incompleteAppDelegate application:OCMClassMock([UIApplication class])
+      didReceiveRemoteNotification:@{}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_appDidReceiveRemoteNotification);
+}
+
+// If the remote notification with fetch handler is NOT implemented, we will force-implement
+// the backup -application:didReceiveRemoteNotification: method
+- (void)testIncompleteAppDelegateRemoteNotificationWithFetchHandlerMethod {
+  IncompleteAppDelegate *incompleteAppDelegate = [[IncompleteAppDelegate alloc] init];
+  [self.mockProxy swizzleAppDelegateMethods:incompleteAppDelegate];
+  SEL remoteNotificationWithFetchHandler =
+  @selector(application:didReceiveRemoteNotification:fetchCompletionHandler:);
+  XCTAssertFalse([incompleteAppDelegate respondsToSelector:remoteNotificationWithFetchHandler]);
+
+  SEL remoteNotification = @selector(application:didReceiveRemoteNotification:);
+  XCTAssertTrue([incompleteAppDelegate respondsToSelector:remoteNotification]);
+}
+
+- (void)testSwizzledAppDelegateRemoteNotificationMethods {
+  FakeAppDelegate *appDelegate = [[FakeAppDelegate alloc] init];
+  [self.mockProxy swizzleAppDelegateMethods:appDelegate];
+  [appDelegate application:OCMClassMock([UIApplication class]) didReceiveRemoteNotification:@{}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_appDidReceiveRemoteNotification);
+  // Verify our original method was called
+  XCTAssertTrue(appDelegate.remoteNotificationMethodWasCalled);
+
+  // Now call the remote notification with handler method
+  [appDelegate application:OCMClassMock([UIApplication class])
+      didReceiveRemoteNotification:@{}
+      fetchCompletionHandler:^(UIBackgroundFetchResult result) {}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_appDidReceiveRemoteNotificationWithHandler);
+  // Verify our original method was called
+  XCTAssertTrue(appDelegate.remoteNotificationWithFetchHandlerWasCalled);
+}
+
+#if __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
+
+- (void)testListeningForDelegateChangesOnInvalidUserNotificationCenter {
+  id randomObject = @"Random Object that is not a User Notification Center";
+  [self.proxy listenForDelegateChangesInUserNotificationCenter:randomObject];
+  XCTAssertFalse(self.proxy.isObservingUserNotificationDelegateChanges);
+}
+
+- (void)testSwizzlingInvalidUserNotificationCenterDelegate {
+  id randomObject = @"Random Object that is not a User Notification Center Delegate";
+  [self.proxy swizzleUserNotificationCenterDelegate:randomObject];
+  XCTAssertFalse(self.proxy.hasSwizzledUserNotificationDelegate);
+}
+
+- (void)testSwizzlingUserNotificationsCenterDelegate {
+  FakeUserNotificationCenterDelegate *delegate = [[FakeUserNotificationCenterDelegate alloc] init];
+  [self.proxy swizzleUserNotificationCenterDelegate:delegate];
+  XCTAssertTrue(self.proxy.hasSwizzledUserNotificationDelegate);
+}
+
+// Use a fake delegate that doesn't actually implement the needed delegate method.
+// Our swizzled method should still be called.
+
+- (void)testIncompleteUserNotificationCenterDelegateMethod {
+  // Early exit if running on pre iOS 10
+  if (![UNNotification class]) {
+    return;
+  }
+  IncompleteUserNotificationCenterDelegate *delegate =
+      [[IncompleteUserNotificationCenterDelegate alloc] init];
+  [self.mockProxy swizzleUserNotificationCenterDelegate:delegate];
+  SEL selector = @selector(userNotificationCenter:willPresentNotification:withCompletionHandler:);
+  XCTAssertTrue([delegate respondsToSelector:selector]);
+  // Invoking delegate method should also invoke our swizzled method
+  // The swizzled method uses the +sharedProxy, which should be
+  // returning our mocked proxy.
+  // Use non-nil, proper classes, otherwise our SDK bails out.
+  [delegate userNotificationCenter:OCMClassMock([UNUserNotificationCenter class])
+           willPresentNotification:[self generateMockNotification]
+             withCompletionHandler:^(NSUInteger options) {}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_willPresentNotificationWithHandler);
+}
+
+// Use an object that does actually implement the needed method. Both should be called.
+- (void)testSwizzledUserNotificationsCenterDelegate {
+  // Early exit if running on pre iOS 10
+  if (![UNNotification class]) {
+    return;
+  }
+  FakeUserNotificationCenterDelegate *delegate = [[FakeUserNotificationCenterDelegate alloc] init];
+  [self.mockProxy swizzleUserNotificationCenterDelegate:delegate];
+  // Invoking delegate method should also invoke our swizzled method
+  // The swizzled method uses the +sharedProxy, which should be
+  // returning our mocked proxy.
+  // Use non-nil, proper classes, otherwise our SDK bails out.
+  [delegate userNotificationCenter:OCMClassMock([UNUserNotificationCenter class])
+           willPresentNotification:[self generateMockNotification]
+             withCompletionHandler:^(NSUInteger options) {}];
+  // Verify our swizzled method was called
+  OCMVerify(FCM_swizzle_willPresentNotificationWithHandler);
+  // Verify our original method was called
+  XCTAssertTrue(delegate.willPresentWasCalled);
+}
+
+- (id)generateMockNotification {
+  // Stub out: notification.request.content.userInfo
+  id mockNotification = OCMClassMock([UNNotification class]);
+  id mockRequest = OCMClassMock([UNNotificationRequest class]);
+  id mockContent = OCMClassMock([UNNotificationContent class]);
+  OCMStub([mockContent userInfo]).andReturn(@{});
+  OCMStub([mockRequest content]).andReturn(mockContent);
+  OCMStub([mockNotification request]).andReturn(mockRequest);
+  return mockNotification;
+}
+
+#endif // __IPHONE_OS_VERSION_MAX_ALLOWED >= __IPHONE_10_0
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingRmqManagerTest.m b/Example/Messaging/Tests/FIRMessagingRmqManagerTest.m
new file mode 100644
index 0000000..88d7c4e
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingRmqManagerTest.m
@@ -0,0 +1,332 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <XCTest/XCTest.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+
+#import "FIRMessagingPersistentSyncMessage.h"
+#import "FIRMessagingRmqManager.h"
+#import "FIRMessagingUtilities.h"
+
+static NSString *const kRmqDatabaseName = @"rmq-test-db";
+static NSString *const kRmqDataMessageCategory = @"com.google.gcm-rmq-test";
+
+@interface FIRMessagingRmqManagerTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) FIRMessagingRmqManager *rmqManager;
+
+@end
+
+@implementation FIRMessagingRmqManagerTest
+
+- (void)setUp {
+  [super setUp];
+  // Make sure we start off with a clean state each time
+  [FIRMessagingRmqManager removeDatabaseWithName:kRmqDatabaseName];
+  _rmqManager = [[FIRMessagingRmqManager alloc] initWithDatabaseName:kRmqDatabaseName];
+}
+
+- (void)tearDown {
+  [super tearDown];
+  [FIRMessagingRmqManager removeDatabaseWithName:kRmqDatabaseName];
+}
+
+/**
+ *  Add s2d messages with different RMQ-ID's to the RMQ. Fetch the messages
+ *  and verify that all messages were successfully saved.
+ */
+- (void)testSavingS2dMessages {
+  NSArray *messageIDs = @[ @"message1", @"message2", @"123456" ];
+  for (NSString *messageID in messageIDs) {
+    [self.rmqManager saveS2dMessageWithRmqId:messageID];
+  }
+  NSArray *rmqMessages = [self.rmqManager unackedS2dRmqIds];
+  XCTAssertEqual(messageIDs.count, rmqMessages.count);
+  for (NSString *messageID in rmqMessages) {
+    XCTAssertTrue([messageIDs containsObject:messageID]);
+  }
+}
+
+/**
+ *  Add s2d messages with different RMQ-ID's to the RMQ. Delete some of the
+ *  messages stored, assuming we received a server ACK for them. The remaining
+ *  messages should be fetched successfully.
+ */
+- (void)testDeletingS2dMessages {
+  NSArray *addMessages = @[ @"message1", @"message2", @"message3", @"message4"];
+  for (NSString *messageID in addMessages) {
+    [self.rmqManager saveS2dMessageWithRmqId:messageID];
+  }
+  NSArray *removeMessages = @[ addMessages[1], addMessages[3] ];
+  [self.rmqManager removeS2dIds:removeMessages];
+  NSArray *remainingMessages = [self.rmqManager unackedS2dRmqIds];
+  XCTAssertEqual(2, remainingMessages.count);
+  XCTAssertTrue([remainingMessages containsObject:addMessages[0]]);
+  XCTAssertTrue([remainingMessages containsObject:addMessages[2]]);
+}
+
+/**
+ *  Test deleting a s2d message that is not in the persistent store. This shouldn't
+ *  crash or alter the valid contents of the RMQ store.
+ */
+- (void)testDeletingInvalidS2dMessage {
+  NSString *validMessageID = @"validMessage123";
+  [self.rmqManager saveS2dMessageWithRmqId:validMessageID];
+  NSString *invalidMessageID = @"invalidMessage123";
+  [self.rmqManager removeS2dIds:@[invalidMessageID]];
+  NSArray *remainingMessages = [self.rmqManager unackedS2dRmqIds];
+  XCTAssertEqual(1, remainingMessages.count);
+  XCTAssertEqualObjects(validMessageID, remainingMessages[0]);
+}
+
+/**
+ *  Test loading the RMQ-ID for d2s messages when there are no outgoing messages in the RMQ.
+ */
+- (void)testLoadRmqIDWithNoD2sMessages {
+  [self.rmqManager loadRmqId];
+  XCTAssertEqual(-1, [self maxRmqIDInRmqStoreForD2SMessages]);
+}
+
+/**
+ *  Test that outgoing RMQ messages are correctly saved
+ */
+- (void)testOutgoingRmqWithValidMessages {
+  NSString *from = @"rmq-test";
+  [self.rmqManager loadRmqId];
+  GtalkDataMessageStanza *message1 = [self dataMessageWithMessageID:@"message1"
+                                                               from:from
+                                                               data:nil];
+  NSError *error;
+
+  // should successfully save the message to RMQ
+  XCTAssertTrue([self.rmqManager saveRmqMessage:message1 error:&error]);
+  XCTAssertNil(error);
+
+  GtalkDataMessageStanza *message2 = [self dataMessageWithMessageID:@"message2"
+                                                               from:from
+                                                               data:nil];
+
+  // should successfully save the second message to RMQ
+  XCTAssertTrue([self.rmqManager saveRmqMessage:message2 error:&error]);
+  XCTAssertNil(error);
+
+  // message1 should have RMQ-ID = 2, message2 = 3
+  XCTAssertEqual(3, [self maxRmqIDInRmqStoreForD2SMessages]);
+  [self.rmqManager scanWithRmqMessageHandler:nil
+                          dataMessageHandler:^(int64_t rmqId, GtalkDataMessageStanza *stanza) {
+                            if (rmqId == 2) {
+                              XCTAssertEqualObjects(@"message1", stanza.id_p);
+                            } else if (rmqId == 3) {
+                              XCTAssertEqualObjects(@"message2", stanza.id_p);
+                            } else {
+                              XCTFail(@"Invalid RmqID %lld for s2d message", rmqId);
+                            }
+                          }];
+}
+
+/**
+ *  Test that an outgoing message with different properties is correctly saved to the RMQ.
+ */
+- (void)testOutgoingDataMessageIsCorrectlySaved {
+  NSString *from = @"rmq-test";
+  NSString *messageID = @"message123";
+  NSString *to = @"to-senderID-123";
+  int32_t ttl = 2400;
+  NSString *registrationToken = @"registration-token";
+  NSDictionary *data = @{
+    @"hello" : @"world",
+    @"count" : @"2",
+  };
+
+  [self.rmqManager loadRmqId];
+  GtalkDataMessageStanza *message = [self dataMessageWithMessageID:messageID
+                                                               from:from
+                                                               data:data];
+  [message setTo:to];
+  [message setTtl:ttl];
+  [message setRegId:registrationToken];
+  NSError *error;
+
+  // should successfully save the message to RMQ
+  XCTAssertTrue([self.rmqManager saveRmqMessage:message error:&error]);
+  XCTAssertNil(error);
+
+  [self.rmqManager scanWithRmqMessageHandler:nil
+                          dataMessageHandler:^(int64_t rmqId, GtalkDataMessageStanza *stanza) {
+                            XCTAssertEqualObjects(from, stanza.from);
+                            XCTAssertEqualObjects(messageID, stanza.id_p);
+                            XCTAssertEqualObjects(to, stanza.to);
+                            XCTAssertEqualObjects(registrationToken, stanza.regId);
+                            XCTAssertEqual(ttl, stanza.ttl);
+                            NSMutableDictionary *d = [NSMutableDictionary dictionary];
+                            for (GtalkAppData *appData in stanza.appDataArray) {
+                              d[appData.key] = appData.value;
+                            }
+                            XCTAssertTrue([data isEqualToDictionary:d]);
+                          }];
+}
+
+/**
+ *  Test D2S messages being deleted from RMQ.
+ */
+- (void)testDeletingD2SMessagesFromRMQ {
+  NSString *message1 = @"message123";
+  NSString *ackedMessage = @"message234";
+  NSString *from = @"from-rmq-test";
+  GtalkDataMessageStanza *stanza1 = [self dataMessageWithMessageID:message1 from:from data:nil];
+  GtalkDataMessageStanza *stanza2 = [self dataMessageWithMessageID:ackedMessage
+                                                              from:from
+                                                              data:nil];
+  NSError *error;
+  XCTAssertTrue([self.rmqManager saveRmqMessage:stanza1 error:&error]);
+  XCTAssertNil(error);
+  XCTAssertTrue([self.rmqManager saveRmqMessage:stanza2 error:&error]);
+  XCTAssertNil(error);
+
+  __block int64_t ackedMessageRmqID = -1;
+  [self.rmqManager scanWithRmqMessageHandler:nil
+                          dataMessageHandler:^(int64_t rmqId, GtalkDataMessageStanza *stanza) {
+                            if ([stanza.id_p isEqualToString:ackedMessage]) {
+                              ackedMessageRmqID = rmqId;
+                            }
+                          }];
+  // should be a valid RMQ ID
+  XCTAssertTrue(ackedMessageRmqID > 0);
+
+  // delete the acked message
+  NSString *rmqIDString = [NSString stringWithFormat:@"%lld", ackedMessageRmqID];
+  XCTAssertEqual(1, [self.rmqManager removeRmqMessagesWithRmqId:rmqIDString]);
+
+  // should only have one message in the d2s RMQ
+  [self.rmqManager scanWithRmqMessageHandler:nil
+                          dataMessageHandler:^(int64_t rmqId, GtalkDataMessageStanza *stanza) {
+                            // the acked message was queued later so should have
+                            // rmqID = ackedMessageRMQID - 1
+                            XCTAssertEqual(ackedMessageRmqID - 1, rmqId);
+                            XCTAssertEqual(message1, stanza2.id_p);
+                          }];
+}
+
+/**
+ *  Test saving a sync message to SYNC_RMQ.
+ */
+- (void)testSavingSyncMessage {
+  NSString *rmqID = @"fake-rmq-id-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 1;
+  XCTAssertTrue([self.rmqManager saveSyncMessageWithRmqID:rmqID
+                                           expirationTime:expirationTime
+                                             apnsReceived:YES
+                                              mcsReceived:NO
+                                                    error:nil]);
+
+  FIRMessagingPersistentSyncMessage *persistentMessage = [self.rmqManager querySyncMessageWithRmqID:rmqID];
+  XCTAssertEqual(persistentMessage.expirationTime, expirationTime);
+  XCTAssertTrue(persistentMessage.apnsReceived);
+  XCTAssertFalse(persistentMessage.mcsReceived);
+}
+
+/**
+ *  Test updating a sync message initially received via MCS, now being received via APNS.
+ */
+- (void)testUpdateMessageReceivedViaAPNS {
+  NSString *rmqID = @"fake-rmq-id-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 1;
+  XCTAssertTrue([self.rmqManager saveSyncMessageWithRmqID:rmqID
+                                           expirationTime:expirationTime
+                                             apnsReceived:NO
+                                              mcsReceived:YES
+                                                    error:nil]);
+
+  // Message was now received via APNS
+  XCTAssertTrue([self.rmqManager updateSyncMessageViaAPNSWithRmqID:rmqID error:nil]);
+
+  FIRMessagingPersistentSyncMessage *persistentMessage = [self.rmqManager querySyncMessageWithRmqID:rmqID];
+  XCTAssertTrue(persistentMessage.apnsReceived);
+  XCTAssertTrue(persistentMessage.mcsReceived);
+}
+
+/**
+ *  Test updating a sync message initially received via APNS, now being received via MCS.
+ */
+- (void)testUpdateMessageReceivedViaMCS {
+  NSString *rmqID = @"fake-rmq-id-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 1;
+  XCTAssertTrue([self.rmqManager saveSyncMessageWithRmqID:rmqID
+                                           expirationTime:expirationTime
+                                             apnsReceived:YES
+                                              mcsReceived:NO
+                                                    error:nil]);
+
+  // Message was now received via APNS
+  XCTAssertTrue([self.rmqManager updateSyncMessageViaMCSWithRmqID:rmqID error:nil]);
+
+  FIRMessagingPersistentSyncMessage *persistentMessage = [self.rmqManager querySyncMessageWithRmqID:rmqID];
+  XCTAssertTrue(persistentMessage.apnsReceived);
+  XCTAssertTrue(persistentMessage.mcsReceived);
+}
+
+/**
+ *  Test deleting sync messages from SYNC_RMQ.
+ */
+- (void)testDeleteSyncMessage {
+  NSString *rmqID = @"fake-rmq-id-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 1;
+  XCTAssertTrue([self.rmqManager saveSyncMessageWithRmqID:rmqID
+                                           expirationTime:expirationTime
+                                             apnsReceived:YES
+                                              mcsReceived:NO
+                                                    error:nil]);
+  XCTAssertNotNil([self.rmqManager querySyncMessageWithRmqID:rmqID]);
+
+  // should successfully delete the message
+  XCTAssertTrue([self.rmqManager deleteSyncMessageWithRmqID:rmqID]);
+  XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:rmqID]);
+}
+
+#pragma mark - Private Helpers
+
+- (GtalkDataMessageStanza *)dataMessageWithMessageID:(NSString *)messageID
+                                                from:(NSString *)from
+                                                data:(NSDictionary *)data {
+  GtalkDataMessageStanza *stanza = [[GtalkDataMessageStanza alloc] init];
+  [stanza setId_p:messageID];
+  [stanza setFrom:from];
+  [stanza setCategory:kRmqDataMessageCategory];
+
+  for (NSString *key in data) {
+    NSString *val = data[key];
+    GtalkAppData *appData = [[GtalkAppData alloc] init];
+    [appData setKey:key];
+    [appData setValue:val];
+    [[stanza appDataArray] addObject:appData];
+  }
+
+  return stanza;
+}
+
+- (int64_t)maxRmqIDInRmqStoreForD2SMessages {
+  __block int64_t maxRmqID = -1;
+  [self.rmqManager scanWithRmqMessageHandler:^(int64_t rmqId, int8_t tag, NSData *data) {
+    if (rmqId > maxRmqID) {
+      maxRmqID = rmqId;
+    }
+  }
+                          dataMessageHandler:nil];
+  return maxRmqID;
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingSecureSocketTest.m b/Example/Messaging/Tests/FIRMessagingSecureSocketTest.m
new file mode 100644
index 0000000..9f6186b
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingSecureSocketTest.m
@@ -0,0 +1,323 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "Protos/GtalkCore.pbobjc.h"
+
+#import "FIRMessagingConnection.h"
+#import "FIRMessagingFakeSocket.h"
+#import "FIRMessagingSecureSocket.h"
+#import "FIRMessagingUtilities.h"
+
+@interface FIRMessagingConnection ()
+
++ (GtalkLoginRequest *)loginRequestWithToken:(NSString *)token authID:(NSString *)authID;
+
+@end
+
+@interface FIRMessagingSecureSocket() <NSStreamDelegate>
+
+@property(nonatomic, readwrite, assign) FIRMessagingSecureSocketState state;
+@property(nonatomic, readwrite, strong) NSInputStream *inStream;
+@property(nonatomic, readwrite, strong) NSOutputStream *outStream;
+
+@property(nonatomic, readwrite, assign) BOOL isVersionSent;
+@property(nonatomic, readwrite, assign) BOOL isVersionReceived;
+@property(nonatomic, readwrite, assign) BOOL isInStreamOpen;
+@property(nonatomic, readwrite, assign) BOOL isOutStreamOpen;
+
+@property(nonatomic, readwrite, strong) NSRunLoop *runLoop;
+
+- (BOOL)performRead;
+
+@end
+
+typedef void(^FIRMessagingTestSocketDisconnectHandler)(void);
+typedef void(^FIRMessagingTestSocketConnectHandler)(void);
+
+@interface FIRMessagingSecureSocketTest : XCTestCase <FIRMessagingSecureSocketDelegate>
+
+@property(nonatomic, readwrite, strong) FIRMessagingFakeSocket *socket;
+@property(nonatomic, readwrite, strong) id mockSocket;
+@property(nonatomic, readwrite, strong) NSError *protoParseError;
+@property(nonatomic, readwrite, strong) GPBMessage *protoReceived;
+@property(nonatomic, readwrite, assign) int8_t protoTagReceived;
+
+@property(nonatomic, readwrite, copy) FIRMessagingTestSocketDisconnectHandler disconnectHandler;
+@property(nonatomic, readwrite, copy) FIRMessagingTestSocketConnectHandler connectHandler;
+
+@end
+
+static BOOL isSafeToDisconnectSocket = NO;
+
+@implementation FIRMessagingSecureSocketTest
+
+- (void)setUp {
+  [super setUp];
+  isSafeToDisconnectSocket = NO;
+  self.protoParseError = nil;
+  self.protoReceived = nil;
+  self.protoTagReceived = 0;
+}
+
+- (void)tearDown {
+  self.disconnectHandler = nil;
+  self.connectHandler = nil;
+  isSafeToDisconnectSocket = YES;
+  [self.socket disconnect];
+  [super tearDown];
+}
+
+#pragma mark - Test Reading
+
+- (void)testSendingVersion {
+  // read as soon as 1 byte is written
+  [self createAndConnectSocketWithBufferSize:1];
+  uint8_t versionByte = 40;
+  [self.socket.outStream write:&versionByte maxLength:1];
+
+  [[[self.mockSocket stub] andDo:^(NSInvocation *invocation) {
+      XCTAssertTrue(isSafeToDisconnectSocket,
+                    @"Should not disconnect socket now");
+  }] disconnect];
+  XCTestExpectation *shouldAcceptVersionExpectation =
+      [self expectationWithDescription:@"Socket should accept version"];
+  dispatch_time_t delay =
+      dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2 * NSEC_PER_SEC));
+  dispatch_after(delay, dispatch_get_main_queue(), ^{
+      XCTAssertTrue(self.socket.isVersionReceived);
+      [shouldAcceptVersionExpectation fulfill];
+  });
+
+  [self waitForExpectationsWithTimeout:3.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+}
+
+- (void)testReceivingDataMessage {
+  [self createAndConnectSocketWithBufferSize:61];
+  [self writeVersionToOutStream];
+  GtalkDataMessageStanza *message = [[GtalkDataMessageStanza alloc] init];
+  [message setCategory:@"socket-test-category"];
+  [message setFrom:@"socket-test-from"];
+  FIRMessagingSetLastStreamId(message, 2);
+  FIRMessagingSetRmq2Id(message, @"socket-test-rmq");
+
+  XCTestExpectation *dataExpectation = [self
+      expectationWithDescription:@"FIRMessaging socket should receive data message"];
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+      [self.socket sendData:[message data]
+                    withTag:kFIRMessagingProtoTagDataMessageStanza
+                      rmqId:FIRMessagingGetRmq2Id(message)];
+  });
+
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(4 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+      XCTAssertEqual(self.protoTagReceived, kFIRMessagingProtoTagDataMessageStanza);
+      [dataExpectation fulfill];
+  });
+
+  [self waitForExpectationsWithTimeout:5.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+}
+
+#pragma mark - Writing
+
+- (void)testLoginRequest {
+  [self createAndConnectSocketWithBufferSize:99];
+
+  XCTestExpectation *loginExpectation =
+      [self expectationWithDescription:@"Socket send valid login proto"];
+  [self writeVersionToOutStream];
+  GtalkLoginRequest *loginRequest =
+      [FIRMessagingConnection loginRequestWithToken:@"gcmtoken" authID:@"gcmauthid"];
+  FIRMessagingSetLastStreamId(loginRequest, 1);
+
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+      [self.socket sendData:[loginRequest data]
+                    withTag:FIRMessagingGetTagForProto(loginRequest)
+                      rmqId:FIRMessagingGetRmq2Id(loginRequest)];
+  });
+
+  dispatch_time_t delay =
+      dispatch_time(DISPATCH_TIME_NOW, (int64_t)(5 * NSEC_PER_SEC));
+  dispatch_after(delay, dispatch_get_main_queue(), ^{
+      XCTAssertTrue(self.socket.isVersionReceived);
+      XCTAssertEqual(self.protoTagReceived, kFIRMessagingProtoTagLoginRequest);
+      [loginExpectation fulfill];
+  });
+
+  [self waitForExpectationsWithTimeout:6.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+}
+
+- (void)testSendingImproperData {
+  [self createAndConnectSocketWithBufferSize:124];
+  [self writeVersionToOutStream];
+
+  NSString *randomString = @"some random data string";
+  NSData *randomData = [randomString dataUsingEncoding:NSUTF8StringEncoding];
+
+  XCTestExpectation *parseErrorExpectation =
+      [self expectationWithDescription:@"Sending improper data results in a parse error"];
+
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+                   [self.socket sendData:randomData withTag:3 rmqId:@"some-random-rmq-id"];
+                 });
+
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2.5 * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+                   if (self.protoParseError != nil) {
+                     [parseErrorExpectation fulfill];
+                   }
+                 });
+
+  [self waitForExpectationsWithTimeout:3.0 handler:nil];
+}
+
+- (void)testSendingDataWithImproperTag {
+  [self createAndConnectSocketWithBufferSize:124];
+  [self writeVersionToOutStream];
+  const char dataString[] = { 0x02, 0x02, 0x11, 0x11, 0x11, 0x11 }; // tag 10, random data
+  NSData *randomData = [NSData dataWithBytes:dataString length:6];
+
+  // Create an expectation for a method which should not be invoked during this test.
+  // This is required to allow us to wait for the socket stream to be read and
+  // processed by FIRMessagingSecureSocket
+  OCMExpect([self.mockSocket disconnect]);
+
+  NSTimeInterval sendDelay = 2.0;
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(sendDelay * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+                   [self.socket sendData:randomData withTag:10 rmqId:@"some-random-rmq-id"];
+                 });
+
+  @try {
+    // While waiting to verify this call, an exception should be thrown
+    // trying to parse the random data in our delegate.
+    // Wait slightly longer than the sendDelay, to allow for the parsing
+    OCMVerifyAllWithDelay(self.mockSocket, sendDelay+0.25);
+    XCTFail(@"Invalid data being read should have thrown an exception.");
+  }
+  @catch (NSException *exception) {
+    XCTAssertNotNil(exception);
+  }
+  @finally { }
+}
+
+- (void)testDisconnect {
+  [self createAndConnectSocketWithBufferSize:1];
+  [self writeVersionToOutStream];
+  // version read and written let's disconnect
+  XCTestExpectation *disconnectExpectation =
+      [self expectationWithDescription:@"socket should disconnect properly"];
+  self.disconnectHandler = ^{
+      [disconnectExpectation fulfill];
+  };
+
+  [self.socket disconnect];
+
+  [self waitForExpectationsWithTimeout:5.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+
+  XCTAssertNil(self.socket.inStream);
+  XCTAssertNil(self.socket.outStream);
+  XCTAssertEqual(self.socket.state, kFIRMessagingSecureSocketClosed);
+}
+
+- (void)testSocketOpening {
+  XCTestExpectation *openSocketExpectation =
+      [self expectationWithDescription:@"Socket should open properly"];
+  self.connectHandler = ^{
+      [openSocketExpectation fulfill];
+  };
+  [self createAndConnectSocketWithBufferSize:1];
+  [self writeVersionToOutStream];
+
+  [self waitForExpectationsWithTimeout:10.0
+                               handler:^(NSError *error) {
+                                   XCTAssertNil(error);
+                               }];
+
+  XCTAssertTrue(self.socket.isInStreamOpen);
+  XCTAssertTrue(self.socket.isOutStreamOpen);
+}
+
+#pragma mark - FIRMessagingSecureSocketDelegate protocol
+
+- (void)secureSocket:(FIRMessagingSecureSocket *)socket
+      didReceiveData:(NSData *)data
+             withTag:(int8_t)tag {
+  NSError *error;
+  GPBMessage *proto =
+      [FIRMessagingGetClassForTag((FIRMessagingProtoTag)tag) parseFromData:data
+                                                                     error:&error];
+  self.protoParseError = error;
+  self.protoReceived = proto;
+  self.protoTagReceived = tag;
+}
+
+- (void)secureSocket:(FIRMessagingSecureSocket *)socket
+ didSendProtoWithTag:(int8_t)tag
+               rmqId:(NSString *)rmqId {
+  // do nothing
+}
+
+- (void)secureSocketDidConnect:(FIRMessagingSecureSocket *)socket {
+  if (self.connectHandler) {
+    self.connectHandler();
+  }
+}
+
+- (void)didDisconnectWithSecureSocket:(FIRMessagingSecureSocket *)socket {
+  if (self.disconnectHandler) {
+    self.disconnectHandler();
+  }
+}
+
+#pragma mark - Private Helpers
+
+- (void)createAndConnectSocketWithBufferSize:(uint8_t)bufferSize {
+  self.socket = [[FIRMessagingFakeSocket alloc] initWithBufferSize:bufferSize];
+  self.mockSocket = OCMPartialMock(self.socket);
+  self.socket.delegate = self;
+
+  [self.socket connectToHost:@"localhost"
+                        port:6234
+                   onRunLoop:[NSRunLoop mainRunLoop]];
+}
+
+- (void)writeVersionToOutStream {
+  uint8_t versionByte = 40;
+  [self.socket.outStream write:&versionByte maxLength:1];
+  // don't resend the version
+  self.socket.isVersionSent = YES;
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingServiceTest.m b/Example/Messaging/Tests/FIRMessagingServiceTest.m
new file mode 100644
index 0000000..bdbc0c2
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingServiceTest.m
@@ -0,0 +1,288 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import UIKit;
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessaging.h"
+#import "FIRMessagingClient.h"
+#import "FIRMessagingConfig.h"
+#import "FIRMessagingPubSub.h"
+#import "FIRMessagingTopicsCommon.h"
+#import "InternalHeaders/FIRMessagingInternalUtilities.h"
+#import "NSError+FIRMessaging.h"
+
+@interface FIRMessaging () <FIRMessagingClientDelegate>
+
+@property(nonatomic, readwrite, strong) FIRMessagingClient *client;
+@property(nonatomic, readwrite, strong) FIRMessagingPubSub *pubsub;
+@property(nonatomic, readwrite, strong) NSString *defaultFcmToken;
+
+@end
+
+@interface FIRMessagingPubSub ()
+
+@property(nonatomic, readwrite, strong) FIRMessagingClient *client;
+
+@end
+
+
+@interface FIRMessagingServiceTest : XCTestCase
+
+@end
+
+@implementation FIRMessagingServiceTest
+
+- (void)testSubscribe {
+  id mockClient = OCMClassMock([FIRMessagingClient class]);
+  FIRMessaging *service = [FIRMessaging messaging];
+  [service setClient:mockClient];
+  [service.pubsub setClient:mockClient];
+
+  XCTestExpectation *subscribeExpectation =
+      [self expectationWithDescription:@"Should call subscribe on FIRMessagingClient"];
+  NSString *token = @"abcdefghijklmn";
+  NSString *topic = @"/topics/some-random-topic";
+
+  [[[mockClient stub]
+    andDo:^(NSInvocation *invocation) {
+      [subscribeExpectation fulfill];
+    }]
+      updateSubscriptionWithToken:token
+                            topic:topic
+                          options:OCMOCK_ANY
+                     shouldDelete:NO
+                          handler:OCMOCK_ANY];
+
+  [service.pubsub subscribeWithToken:token
+                               topic:topic
+                             options:nil
+                             handler:^(FIRMessagingTopicOperationResult result, NSError *error) {
+                               // not a nil block
+                             }];
+
+  // should call updateSubscription
+  [self waitForExpectationsWithTimeout:0.1
+                               handler:^(NSError *error) {
+                                 XCTAssertNil(error);
+                                 [mockClient verify];
+                               }];
+}
+
+- (void)testUnsubscribe {
+  id mockClient = OCMClassMock([FIRMessagingClient class]);
+  FIRMessaging *messaging = [FIRMessaging messaging];
+  [messaging setClient:mockClient];
+  [messaging.pubsub setClient:mockClient];
+
+  XCTestExpectation *subscribeExpectation =
+      [self expectationWithDescription:@"Should call unsubscribe on FIRMessagingClient"];
+
+  NSString *token = @"abcdefghijklmn";
+  NSString *topic = @"/topics/some-random-topic";
+
+  [[[mockClient stub] andDo:^(NSInvocation *invocation) {
+    [subscribeExpectation fulfill];
+  }]
+      updateSubscriptionWithToken:[OCMArg isEqual:token]
+                            topic:[OCMArg isEqual:topic]
+                          options:[OCMArg checkWithBlock:^BOOL(id obj) {
+                            if ([obj isKindOfClass:[NSDictionary class]]) {
+                              return [(NSDictionary *)obj count] == 0;
+                            }
+                            return NO;
+                          }]
+                     shouldDelete:YES
+                          handler:OCMOCK_ANY];
+
+  [messaging.pubsub unsubscribeWithToken:token
+                                   topic:topic
+                                 options:nil
+                                 handler:^(FIRMessagingTopicOperationResult result, NSError *error){
+
+                                 }];
+
+  // should call updateSubscription
+  [self waitForExpectationsWithTimeout:0.1
+                               handler:^(NSError *error) {
+                                 XCTAssertNil(error);
+                                 [mockClient verify];
+                               }];
+}
+
+/**
+ *  Test using PubSub without explicitly starting FIRMessagingService.
+ */
+- (void)testSubscribeWithoutStart {
+  [[[FIRMessaging messaging] pubsub] subscribeWithToken:@"abcdef1234"
+                                                  topic:@"/topics/hello-world"
+                                                options:nil
+                                                handler:
+      ^(FIRMessagingTopicOperationResult result, NSError *error) {
+    XCTAssertNil(error);
+    XCTAssertEqual(kFIRMessagingErrorCodePubSubFIRMessagingNotSetup,
+                   error.code);
+  }];
+}
+
+- (void)testSubscribeWithInvalidToken {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+
+  XCTestExpectation *exceptionExpectation =
+      [self expectationWithDescription:@"Should throw exception for invalid token"];
+  @try {
+    [messaging.pubsub subscribeWithToken:@""
+                                   topic:@"/topics/hello-world"
+                                 options:nil
+                                 handler:
+        ^(FIRMessagingTopicOperationResult result, NSError *error) {
+      XCTFail(@"Should not invoke the handler");
+    }];
+  }
+  @catch (NSException *exception) {
+    [exceptionExpectation fulfill];
+  }
+  @finally {
+    [self waitForExpectationsWithTimeout:0.1 handler:^(NSError *error) {
+      XCTAssertNil(error);
+    }];
+  }
+}
+
+- (void)testUnsubscribeWithInvalidTopic {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+
+  XCTestExpectation *exceptionExpectation =
+      [self expectationWithDescription:@"Should throw exception for invalid token"];
+  @try {
+    [messaging.pubsub unsubscribeWithToken:@"abcdef1234"
+                                     topic:nil
+                                   options:nil
+                                   handler:
+        ^(FIRMessagingTopicOperationResult result, NSError *error) {
+      XCTFail(@"Should not invoke the handler");
+    }];
+  }
+  @catch (NSException *exception) {
+    [exceptionExpectation fulfill];
+  }
+  @finally {
+    [self waitForExpectationsWithTimeout:0.1 handler:^(NSError *error) {
+      XCTAssertNil(error);
+    }];
+  }
+}
+
+- (void)testSubscribeWithNoTopicPrefix {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+  FIRMessagingPubSub *pubSub = messaging.pubsub;
+  id mockPubSub = OCMClassMock([FIRMessagingPubSub class]);
+
+  NSString *topicName = @"topicWithoutPrefix";
+  NSString *topicNameWithPrefix = [FIRMessagingPubSub addPrefixToTopic:topicName];
+  messaging.pubsub = mockPubSub;
+  messaging.defaultFcmToken = @"fake-default-token";
+  OCMExpect([messaging.pubsub subscribeToTopic:[OCMArg isEqual:topicNameWithPrefix]]);
+  [messaging subscribeToTopic:topicName];
+  OCMVerifyAll(mockPubSub);
+  // Need to swap back since it's a singleton and hence will live beyond the scope of this test.
+  messaging.pubsub = pubSub;
+}
+
+- (void)testSubscribeWithTopicPrefix {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+  FIRMessagingPubSub *pubSub = messaging.pubsub;
+  id mockPubSub = OCMClassMock([FIRMessagingPubSub class]);
+
+  NSString *topicName = @"/topics/topicWithoutPrefix";
+  messaging.pubsub = mockPubSub;
+  messaging.defaultFcmToken = @"fake-default-token";
+  OCMExpect([messaging.pubsub subscribeToTopic:[OCMArg isEqual:topicName]]);
+  [messaging subscribeToTopic:topicName];
+  OCMVerifyAll(mockPubSub);
+  // Need to swap back since it's a singleton and hence will live beyond the scope of this test.
+  messaging.pubsub = pubSub;
+}
+
+- (void)testUnsubscribeWithNoTopicPrefix {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+  FIRMessagingPubSub *pubSub = messaging.pubsub;
+  id mockPubSub = OCMClassMock([FIRMessagingPubSub class]);
+
+  NSString *topicName = @"topicWithoutPrefix";
+  NSString *topicNameWithPrefix = [FIRMessagingPubSub addPrefixToTopic:topicName];
+  messaging.pubsub = mockPubSub;
+  messaging.defaultFcmToken = @"fake-default-token";
+  OCMExpect([messaging.pubsub unsubscribeFromTopic:[OCMArg isEqual:topicNameWithPrefix]]);
+  [messaging unsubscribeFromTopic:topicName];
+  OCMVerifyAll(mockPubSub);
+  // Need to swap back since it's a singleton and hence will live beyond the scope of this test.
+  messaging.pubsub = pubSub;
+}
+
+- (void)testUnsubscribeWithTopicPrefix {
+  FIRMessaging *messaging = [FIRMessaging messaging];
+  FIRMessagingPubSub *pubSub = messaging.pubsub;
+  id mockPubSub = OCMClassMock([FIRMessagingPubSub class]);
+
+  NSString *topicName = @"/topics/topicWithPrefix";
+  messaging.pubsub = mockPubSub;
+  messaging.defaultFcmToken = @"fake-default-token";
+  OCMExpect([messaging.pubsub unsubscribeFromTopic:[OCMArg isEqual:topicName]]);
+  [messaging unsubscribeFromTopic:topicName];
+  OCMVerifyAll(mockPubSub);
+  // Need to swap back since it's a singleton and hence will live beyond the scope of this test.
+  messaging.pubsub = pubSub;
+}
+
+- (void)testFIRMessagingSDKVersionInFIRMessagingService {
+  Class versionClass = NSClassFromString(kFIRMessagingSDKClassString);
+  SEL versionSelector = NSSelectorFromString(kFIRMessagingSDKVersionSelectorString);
+  if ([versionClass respondsToSelector:versionSelector]) {
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
+    id versionString = [versionClass performSelector:versionSelector];
+#pragma clang diagnostic pop
+
+    XCTAssertTrue([versionString isKindOfClass:[NSString class]]);
+  } else {
+    XCTFail("%@ does not respond to selector %@",
+            kFIRMessagingSDKClassString, kFIRMessagingSDKVersionSelectorString);
+  }
+}
+
+- (void)testFIRMessagingSDKLocaleInFIRMessagingService {
+  Class klass = NSClassFromString(kFIRMessagingSDKClassString);
+  SEL localeSelector = NSSelectorFromString(kFIRMessagingSDKLocaleSelectorString);
+  if ([klass respondsToSelector:localeSelector]) {
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Warc-performSelector-leaks"
+    id locale = [klass performSelector:localeSelector];
+#pragma clang diagnostic pop
+
+    XCTAssertTrue([locale isKindOfClass:[NSString class]]);
+    XCTAssertNotNil(locale);
+  } else {
+    XCTFail("%@ does not respond to selector %@",
+            kFIRMessagingSDKClassString, kFIRMessagingSDKLocaleSelectorString);
+  }
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingSyncMessageManagerTest.m b/Example/Messaging/Tests/FIRMessagingSyncMessageManagerTest.m
new file mode 100644
index 0000000..e40e1f2
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingSyncMessageManagerTest.m
@@ -0,0 +1,256 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <XCTest/XCTest.h>
+
+#import "FIRMessagingPersistentSyncMessage.h"
+#import "FIRMessagingRmqManager.h"
+#import "FIRMessagingSyncMessageManager.h"
+#import "FIRMessagingUtilities.h"
+#import "FIRMessagingConstants.h"
+
+static NSString *const kRmqSqliteFilename = @"rmq-sync-manager-test";
+
+@interface FIRMessagingSyncMessageManagerTest : XCTestCase
+
+@property(nonatomic, readwrite, strong) FIRMessagingRmqManager *rmqManager;
+@property(nonatomic, readwrite, strong) FIRMessagingSyncMessageManager *syncMessageManager;
+
+@end
+
+@implementation FIRMessagingSyncMessageManagerTest
+
+- (void)setUp {
+  [super setUp];
+  // Make sure the db state is clean before we begin.
+  [FIRMessagingRmqManager removeDatabaseWithName:kRmqSqliteFilename];
+  self.rmqManager = [[FIRMessagingRmqManager alloc] initWithDatabaseName:kRmqSqliteFilename];
+  self.syncMessageManager = [[FIRMessagingSyncMessageManager alloc] initWithRmqManager:self.rmqManager];
+}
+
+- (void)tearDown {
+  [[self.rmqManager class] removeDatabaseWithName:kRmqSqliteFilename];
+  [super tearDown];
+}
+
+/**
+ *  Test receiving a new sync message via APNS should be added to SYNC_RMQ.
+ */
+- (void)testNewAPNSMessage {
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 86400;  // 1 day in future
+
+  NSDictionary *oldMessage = @{
+    kFIRMessagingMessageIDKey : @"fake-rmq-1",
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:oldMessage]);
+
+  NSDictionary *newMessage = @{
+    kFIRMessagingMessageIDKey : @"fake-rmq-2",
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:newMessage]);
+}
+
+/**
+ *  Test receiving a new sync message via MCS should be added to SYNC_RMQ.
+ */
+- (void)testNewMCSMessage {
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 86400;  // 1 day in future
+  NSDictionary *oldMessage = @{
+    kFIRMessagingMessageIDKey : @"fake-rmq-1",
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveMCSSyncMessage:oldMessage]);
+
+  NSDictionary *newMessage = @{
+    kFIRMessagingMessageIDKey : @"fake-rmq-2",
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:newMessage]);
+}
+
+/**
+ *  Test receiving a duplicate message via APNS.
+ */
+- (void)testDuplicateAPNSMessage {
+  NSString *messageID = @"fake-rmq-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 86400;  // 1 day in future
+  NSDictionary *newMessage = @{
+    kFIRMessagingMessageIDKey : messageID,
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:newMessage]);
+
+  // The message is a duplicate
+  XCTAssertTrue([self.syncMessageManager didReceiveAPNSSyncMessage:newMessage]);
+
+  FIRMessagingPersistentSyncMessage *persistentMessage =
+      [self.rmqManager querySyncMessageWithRmqID:messageID];
+  XCTAssertTrue(persistentMessage.apnsReceived);
+  XCTAssertFalse(persistentMessage.mcsReceived);
+}
+
+/**
+ *  Test receiving a duplicate message via MCS.
+ */
+- (void)testDuplicateMCSMessage {
+  NSString *messageID = @"fake-rmq-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 86400;  // 1 day in future
+  NSDictionary *newMessage = @{
+    kFIRMessagingMessageIDKey : messageID,
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+
+  XCTAssertFalse([self.syncMessageManager didReceiveMCSSyncMessage:newMessage]);
+
+  // The message is a duplicate
+  XCTAssertTrue([self.syncMessageManager didReceiveMCSSyncMessage:newMessage]);
+
+  FIRMessagingPersistentSyncMessage *persistentMessage =
+      [self.rmqManager querySyncMessageWithRmqID:messageID];
+  XCTAssertFalse(persistentMessage.apnsReceived);
+  XCTAssertTrue(persistentMessage.mcsReceived);
+}
+
+/**
+ *  Test receiving a sync message both via APNS and MCS.
+ */
+- (void)testMessageReceivedBothViaAPNSAndMCS {
+  NSString *messageID = @"fake-rmq-1";
+  int64_t expirationTime = FIRMessagingCurrentTimestampInSeconds() + 86400;  // 1 day in future
+  NSDictionary *newMessage = @{
+    kFIRMessagingMessageIDKey : messageID,
+    kFIRMessagingMessageSyncViaMCSKey : @(expirationTime),
+    @"hello" : @"world",
+  };
+
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:newMessage]);
+  // Duplicate of the above received APNS message
+  XCTAssertTrue([self.syncMessageManager didReceiveMCSSyncMessage:newMessage]);
+
+  // Since we've received both APNS and MCS messages we should have deleted them from SYNC_RMQ
+  FIRMessagingPersistentSyncMessage *persistentMessage =
+      [self.rmqManager querySyncMessageWithRmqID:messageID];
+  XCTAssertNil(persistentMessage);
+}
+
+- (void)testDeletingExpiredMessages {
+  NSString *unexpiredMessageID = @"fake-not-expired-rmqID";
+  int64_t futureExpirationTime = 86400;  // 1 day in future
+  NSDictionary *unexpiredMessage = @{
+    kFIRMessagingMessageIDKey : unexpiredMessageID,
+    kFIRMessagingMessageSyncMessageTTLKey : @(futureExpirationTime),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:unexpiredMessage]);
+
+  NSString *expiredMessageID = @"fake-expired-rmqID";
+  int64_t past = -86400;  // 1 day in past
+  NSDictionary *expiredMessage = @{
+    kFIRMessagingMessageIDKey : expiredMessageID,
+    kFIRMessagingMessageSyncMessageTTLKey : @(past),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:expiredMessage]);
+
+  NSString *noTTLMessageID = @"no-ttl-rmqID"; // no TTL specified should be 4 weeks
+  NSDictionary *noTTLMessage = @{
+    kFIRMessagingMessageIDKey : noTTLMessageID,
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:noTTLMessage]);
+
+  [self.syncMessageManager removeExpiredSyncMessages];
+
+  XCTAssertNotNil([self.rmqManager querySyncMessageWithRmqID:unexpiredMessageID]);
+  XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:expiredMessageID]);
+  XCTAssertNotNil([self.rmqManager querySyncMessageWithRmqID:noTTLMessageID]);
+}
+
+- (void)testDeleteFinishedMessages {
+  NSString *unexpiredMessageID = @"fake-not-expired-rmqID";
+  int64_t futureExpirationTime = 86400;  // 1 day in future
+  NSDictionary *unexpiredMessage = @{
+    kFIRMessagingMessageIDKey : unexpiredMessageID,
+    kFIRMessagingMessageSyncMessageTTLKey : @(futureExpirationTime),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:unexpiredMessage]);
+
+  NSString *noTTLMessageID = @"no-ttl-rmqID"; // no TTL specified should be 4 weeks
+  NSDictionary *noTTLMessage = @{
+                                 kFIRMessagingMessageIDKey : noTTLMessageID,
+                                 @"hello" : @"world",
+                                 };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:noTTLMessage]);
+
+  // Mark the no-TTL message as received via MCS too
+  XCTAssertTrue([self.rmqManager updateSyncMessageViaMCSWithRmqID:noTTLMessageID error:nil]);
+
+  [self.syncMessageManager removeExpiredSyncMessages];
+
+  XCTAssertNotNil([self.rmqManager querySyncMessageWithRmqID:unexpiredMessageID]);
+  XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:noTTLMessageID]);
+}
+
+- (void)testDeleteFinishedAndExpiredMessages {
+  NSString *unexpiredMessageID = @"fake-not-expired-rmqID";
+  int64_t futureExpirationTime = 86400;  // 1 day in future
+  NSDictionary *unexpiredMessage = @{
+    kFIRMessagingMessageIDKey : unexpiredMessageID,
+    kFIRMessagingMessageSyncMessageTTLKey : @(futureExpirationTime),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:unexpiredMessage]);
+
+  NSString *expiredMessageID = @"fake-expired-rmqID";
+  int64_t past = -86400;  // 1 day in past
+  NSDictionary *expiredMessage = @{
+    kFIRMessagingMessageIDKey : expiredMessageID,
+    kFIRMessagingMessageSyncMessageTTLKey : @(past),
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:expiredMessage]);
+
+  NSString *noTTLMessageID = @"no-ttl-rmqID"; // no TTL specified should be 4 weeks
+  NSDictionary *noTTLMessage = @{
+    kFIRMessagingMessageIDKey : noTTLMessageID,
+    @"hello" : @"world",
+  };
+  XCTAssertFalse([self.syncMessageManager didReceiveAPNSSyncMessage:noTTLMessage]);
+
+  // Mark the no-TTL message as received via MCS too
+  XCTAssertTrue([self.rmqManager updateSyncMessageViaMCSWithRmqID:noTTLMessageID error:nil]);
+
+  // Remove expired or finished sync messages.
+  [self.syncMessageManager removeExpiredSyncMessages];
+
+  XCTAssertNotNil([self.rmqManager querySyncMessageWithRmqID:unexpiredMessageID]);
+  XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:expiredMessageID]);
+  XCTAssertNil([self.rmqManager querySyncMessageWithRmqID:noTTLMessageID]);
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingTest.m b/Example/Messaging/Tests/FIRMessagingTest.m
new file mode 100644
index 0000000..3d7c95f
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingTest.m
@@ -0,0 +1,214 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@import XCTest;
+
+#import <OCMock/OCMock.h>
+
+#import "FIRMessaging.h"
+#import "FIRMessagingConfig.h"
+#import "FIRMessagingInstanceIDProxy.h"
+
+extern NSString *const kFIRMessagingFCMTokenFetchAPNSOption;
+
+@interface FIRMessaging ()
+
+@property(nonatomic, readwrite, strong) NSString *defaultFcmToken;
+@property(nonatomic, readwrite, strong) NSData *apnsTokenData;
+@property(nonatomic, readwrite, strong) FIRMessagingInstanceIDProxy *instanceIDProxy;
+
+- (instancetype)initWithConfig:(FIRMessagingConfig *)config;
+// Direct Channel Methods
+- (void)updateAutomaticClientConnection;
+- (BOOL)shouldBeConnectedAutomatically;
+
+@end
+
+@interface FIRMessagingTest : XCTestCase
+
+@property(nonatomic, readonly, strong) FIRMessaging *messaging;
+@property(nonatomic, readwrite, strong) id mockMessaging;
+@property(nonatomic, readwrite, strong) id mockInstanceIDProxy;
+
+@end
+
+@implementation FIRMessagingTest
+
+- (void)setUp {
+  [super setUp];
+  FIRMessagingConfig *config = [FIRMessagingConfig defaultConfig];
+  _messaging = [[FIRMessaging alloc] initWithConfig:config];
+  _mockMessaging = OCMPartialMock(self.messaging);
+  _mockInstanceIDProxy = OCMPartialMock(self.messaging.instanceIDProxy);
+  self.messaging.instanceIDProxy = _mockInstanceIDProxy;
+}
+
+- (void)tearDown {
+  _messaging = nil;
+  _mockMessaging = nil;
+  [super tearDown];
+}
+
+#pragma mark - Direct Channel Establishment Testing
+
+// Should connect with valid token and application in foreground
+- (void)testDoesAutomaticallyConnectIfTokenAvailableAndForegrounded {
+  // Disable actually attempting a connection
+  [[[_mockMessaging stub] andDo:^(NSInvocation *invocation) {
+    // Doing nothing on purpose, when -updateAutomaticClientConnection is called
+  }] updateAutomaticClientConnection];
+  // Set direct channel to be established after disabling connection attempt
+  self.messaging.shouldEstablishDirectChannel = YES;
+  // Set a "valid" token (i.e. not nil or empty)
+  self.messaging.defaultFcmToken = @"1234567";
+  // Swizzle application state to return UIApplicationStateActive
+  UIApplication *app = [UIApplication sharedApplication];
+  id mockApp = OCMPartialMock(app);
+  [[[mockApp stub] andReturnValue:@(UIApplicationStateActive)] applicationState];
+  BOOL shouldBeConnected = [_mockMessaging shouldBeConnectedAutomatically];
+  XCTAssertTrue(shouldBeConnected);
+}
+
+// Should not connect if application is active, but token is empty
+- (void)testDoesNotAutomaticallyConnectIfTokenIsEmpty {
+  // Disable actually attempting a connection
+  [[[_mockMessaging stub] andDo:^(NSInvocation *invocation) {
+    // Doing nothing on purpose, when -updateAutomaticClientConnection is called
+  }] updateAutomaticClientConnection];
+  // Set direct channel to be established after disabling connection attempt
+  self.messaging.shouldEstablishDirectChannel = YES;
+  // By default, there should be no fcmToken
+  // Swizzle application state to return UIApplicationStateActive
+  UIApplication *app = [UIApplication sharedApplication];
+  id mockApp = OCMPartialMock(app);
+  [[[mockApp stub] andReturnValue:@(UIApplicationStateActive)] applicationState];
+  BOOL shouldBeConnected = [_mockMessaging shouldBeConnectedAutomatically];
+  XCTAssertFalse(shouldBeConnected);
+}
+
+// Should not connect if token valid but application isn't active
+- (void)testDoesNotAutomaticallyConnectIfApplicationNotActive {
+  // Disable actually attempting a connection
+  [[[_mockMessaging stub] andDo:^(NSInvocation *invocation) {
+    // Doing nothing on purpose, when -updateAutomaticClientConnection is called
+  }] updateAutomaticClientConnection];
+  // Set direct channel to be established after disabling connection attempt
+  self.messaging.shouldEstablishDirectChannel = YES;
+  // Set a "valid" token (i.e. not nil or empty)
+  self.messaging.defaultFcmToken = @"abcd1234";
+  // Swizzle application state to return UIApplicationStateActive
+  UIApplication *app = [UIApplication sharedApplication];
+  id mockApp = OCMPartialMock(app);
+  [[[mockApp stub] andReturnValue:@(UIApplicationStateBackground)] applicationState];
+  BOOL shouldBeConnected = [_mockMessaging shouldBeConnectedAutomatically];
+  XCTAssertFalse(shouldBeConnected);
+}
+
+#pragma mark - FCM Token Fetching and Deleting
+
+#ifdef NEED_WORKAROUND_FOR_PRIVATE_OCMOCK_getArgumentAtIndexAsObject
+- (void)testAPNSTokenIncludedInOptionsIfAvailableDuringTokenFetch {
+  self.messaging.apnsTokenData =
+      [@"PRETENDING_TO_BE_A_DEVICE_TOKEN" dataUsingEncoding:NSUTF8StringEncoding];
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"Included APNS Token data in options dict."];
+  // Inspect the 'options' dictionary to tell whether our expectation was fulfilled
+  [[[self.mockInstanceIDProxy stub] andDo:^(NSInvocation *invocation) {
+    // Calling getArgument:atIndex: directly leads to an EXC_BAD_ACCESS; use OCMock's wrapper.
+    NSDictionary *options = [invocation getArgumentAtIndexAsObject:4];
+    if (options[@"apns_token"] != nil) {
+      [expectation fulfill];
+    }
+  }] tokenWithAuthorizedEntity:OCMOCK_ANY scope:OCMOCK_ANY options:OCMOCK_ANY handler:OCMOCK_ANY];
+  [self.messaging retrieveFCMTokenForSenderID:@"123456"
+                                   completion:^(NSString * _Nullable FCMToken,
+                                                NSError * _Nullable error) {}];
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+- (void)testAPNSTokenNotIncludedIfUnavailableDuringTokenFetch {
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"Included APNS Token data not included in options dict."];
+  // Inspect the 'options' dictionary to tell whether our expectation was fulfilled
+  [[[self.mockInstanceIDProxy stub] andDo:^(NSInvocation *invocation) {
+    // Calling getArgument:atIndex: directly leads to an EXC_BAD_ACCESS; use OCMock's wrapper.
+    NSDictionary *options = [invocation getArgumentAtIndexAsObject:4];
+    if (options[@"apns_token"] == nil) {
+      [expectation fulfill];
+    }
+  }] tokenWithAuthorizedEntity:OCMOCK_ANY scope:OCMOCK_ANY options:OCMOCK_ANY handler:OCMOCK_ANY];
+  [self.messaging retrieveFCMTokenForSenderID:@"123456"
+                                   completion:^(NSString * _Nullable FCMToken,
+                                                NSError * _Nullable error) {}];
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+#endif
+
+- (void)testReturnsErrorWhenFetchingTokenWithoutSenderID {
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"Returned an error fetching token without Sender ID"];
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wnonnull"
+  [self.messaging retrieveFCMTokenForSenderID:nil
+                                  completion:
+      ^(NSString * _Nullable FCMToken, NSError * _Nullable error) {
+    if (error != nil) {
+      [expectation fulfill];
+    }
+  }];
+#pragma clang diagnostic pop
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+- (void)testReturnsErrorWhenFetchingTokenWithEmptySenderID {
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"Returned an error fetching token with empty Sender ID"];
+  [self.messaging retrieveFCMTokenForSenderID:@""
+                                  completion:
+      ^(NSString * _Nullable FCMToken, NSError * _Nullable error) {
+    if (error != nil) {
+      [expectation fulfill];
+    }
+  }];
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+- (void)testReturnsErrorWhenDeletingTokenWithoutSenderID {
+  XCTestExpectation *expectation =
+      [self expectationWithDescription:@"Returned an error deleting token without Sender ID"];
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wnonnull"
+  [self.messaging deleteFCMTokenForSenderID:nil completion:^(NSError * _Nullable error) {
+    if (error != nil) {
+      [expectation fulfill];
+    }
+  }];
+#pragma clang diagnostic pop
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+- (void)testReturnsErrorWhenDeletingTokenWithEmptySenderID {
+  XCTestExpectation *expectation =
+  [self expectationWithDescription:@"Returned an error deleting token with empty Sender ID"];
+  [self.messaging deleteFCMTokenForSenderID:@"" completion:^(NSError * _Nullable error) {
+    if (error != nil) {
+      [expectation fulfill];
+    }
+  }];
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.h b/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.h
new file mode 100644
index 0000000..0bc6010
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.h
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+@interface FIRMessagingTestNotificationUtilities : NSObject
+
++ (NSMutableDictionary *)createBasicNotificationWithUniqueMessageID;
+
+@end
diff --git a/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.m b/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.m
new file mode 100644
index 0000000..43842ee
--- /dev/null
+++ b/Example/Messaging/Tests/FIRMessagingTestNotificationUtilities.m
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FIRMessagingTestNotificationUtilities.h"
+
+#import "FIRMessagingConstants.h"
+
+@implementation FIRMessagingTestNotificationUtilities
+
++ (NSMutableDictionary *)createBasicNotificationWithUniqueMessageID {
+  NSMutableDictionary *notification = [NSMutableDictionary dictionary];
+  // Always generate a unique message id
+  notification[kFIRMessagingMessageIDKey] =
+  [NSString stringWithFormat:@"%@", @([NSDate date].timeIntervalSince1970)];
+  return notification;
+}
+
+@end
diff --git a/Example/Messaging/Tests/Info.plist b/Example/Messaging/Tests/Info.plist
new file mode 100644
index 0000000..ba72822
--- /dev/null
+++ b/Example/Messaging/Tests/Info.plist
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>CFBundleDevelopmentRegion</key>
+	<string>en</string>
+	<key>CFBundleExecutable</key>
+	<string>$(EXECUTABLE_NAME)</string>
+	<key>CFBundleIdentifier</key>
+	<string>$(PRODUCT_BUNDLE_IDENTIFIER)</string>
+	<key>CFBundleInfoDictionaryVersion</key>
+	<string>6.0</string>
+	<key>CFBundleName</key>
+	<string>$(PRODUCT_NAME)</string>
+	<key>CFBundlePackageType</key>
+	<string>BNDL</string>
+	<key>CFBundleShortVersionString</key>
+	<string>1.0</string>
+	<key>CFBundleSignature</key>
+	<string>????</string>
+	<key>CFBundleVersion</key>
+	<string>1</string>
+</dict>
+</plist>