diff --git a/ChatSecureTests/OTROMEMOIntegrationTest.swift b/ChatSecureTests/OTROMEMOIntegrationTest.swift
index 14f006fcd..9329d7d49 100644
--- a/ChatSecureTests/OTROMEMOIntegrationTest.swift
+++ b/ChatSecureTests/OTROMEMOIntegrationTest.swift
@@ -20,9 +20,10 @@ class OTROMEMOIntegrationTest: XCTestCase {
     
     // This is teh 'remote' user in this setup
     var aliceUser:TestUser?
+    var aliceOmemoModule:OTROMEMOTestModule?
     // This is the 'local' user
     var bobUser:TestUser?
-    var omemoModule:OTROMEMOTestModule?
+    var bobOmemoModule:OTROMEMOTestModule?
     
     override func setUp() {
         super.setUp()
@@ -35,11 +36,19 @@ class OTROMEMOIntegrationTest: XCTestCase {
         let aliceName = "\(name)-alice"
         let bobName = "\(name)-bob"
         self.aliceUser = self.setupUserWithName(aliceName,buddyName: bobName)
+        self.aliceOmemoModule = OTROMEMOTestModule(OMEMOStorage: self.aliceUser!.signalOMEMOCoordinator, xmlNamespace: .ConversationsLegacy, dispatchQueue: nil)
+        self.aliceOmemoModule?.addDelegate(self.aliceUser!.signalOMEMOCoordinator, delegateQueue: self.aliceUser!.signalOMEMOCoordinator.workQueue)
+        self.aliceOmemoModule?.thisUser = aliceUser
+        
         self.bobUser = self.setupUserWithName(bobName,buddyName: aliceName)
-        self.omemoModule = OTROMEMOTestModule(OMEMOStorage: self.bobUser!.signalOMEMOCoordinator, xmlNamespace: .ConversationsLegacy, dispatchQueue: nil)
-        self.omemoModule?.addDelegate(self.bobUser!.signalOMEMOCoordinator, delegateQueue: self.bobUser!.signalOMEMOCoordinator.workQueue)
-        self.omemoModule?.otherUser = aliceUser
-        self.omemoModule?.thisUser = bobUser
+        self.bobOmemoModule = OTROMEMOTestModule(OMEMOStorage: self.bobUser!.signalOMEMOCoordinator, xmlNamespace: .ConversationsLegacy, dispatchQueue: nil)
+        self.bobOmemoModule?.addDelegate(self.bobUser!.signalOMEMOCoordinator, delegateQueue: self.bobUser!.signalOMEMOCoordinator.workQueue)
+        self.bobOmemoModule?.thisUser = bobUser
+        
+        
+        self.aliceOmemoModule?.otherUser = self.bobOmemoModule
+        self.bobOmemoModule?.otherUser = self.aliceOmemoModule
+        
     }
     
     /** 
@@ -73,7 +82,7 @@ class OTROMEMOIntegrationTest: XCTestCase {
      */
     func testDeviceSetup() {
         self.setupTwoAccounts(#function)
-        self.omemoModule?.xmppStreamDidAuthenticate(nil)
+        self.bobOmemoModule?.xmppStreamDidAuthenticate(nil)
         let buddy = self.bobUser!.buddy
         let connection = self.bobUser?.databaseManager.readOnlyDatabaseConnection
         connection?.readWithBlock({ (transaction) in
@@ -91,7 +100,7 @@ class OTROMEMOIntegrationTest: XCTestCase {
     */
     func testFetchingBundleSetup() {
         self.setupTwoAccounts(#function)
-        self.omemoModule?.xmppStreamDidAuthenticate(nil)
+        self.bobOmemoModule?.xmppStreamDidAuthenticate(nil)
         let expectation = self.expectationWithDescription("Sending Message")
         let messageText = "This is message from Bob to Alice"
         self.bobUser!.signalOMEMOCoordinator.encryptAndSendMessage(messageText, buddyYapKey: self.bobUser!.buddy.uniqueId, messageId: "message1") { (success, error) in
@@ -107,9 +116,11 @@ class OTROMEMOIntegrationTest: XCTestCase {
         var messageFound = false
         self.aliceUser?.databaseManager.readOnlyDatabaseConnection.readWithBlock({ (transaction) in
             transaction.enumerateKeysAndObjectsInCollection(OTRBaseMessage.collection(), usingBlock: { (key, object, stop) in
-                let message = object as! OTRBaseMessage
-                XCTAssertEqual(message.text, messageText)
-                messageFound = true
+                if let message = object as? OTRBaseMessage {
+                    XCTAssertEqual(message.text, messageText)
+                    messageFound = true
+                }
+                
             })
             
             transaction.enumerateKeysAndObjectsInCollection(OTROMEMODevice.collection(), usingBlock: { (key, object, stop) in
@@ -123,7 +134,7 @@ class OTROMEMOIntegrationTest: XCTestCase {
     
     func testRemoveDevice() {
         self.setupTwoAccounts(#function)
-        self.omemoModule?.xmppStreamDidAuthenticate(nil)
+        self.bobOmemoModule?.xmppStreamDidAuthenticate(nil)
         let expectation = self.expectationWithDescription("Remove Devices")
         let deviceNumber = NSNumber(int:5)
         let device = OTROMEMODevice(deviceId: deviceNumber, trustLevel: OMEMOTrustLevel.TrustedTofu, parentKey: self.bobUser!.account.uniqueId, parentCollection: OTRAccount.collection(), publicIdentityKeyData: nil, lastSeenDate: nil)
diff --git a/ChatSecureTests/OTROMEMOTestModule.swift b/ChatSecureTests/OTROMEMOTestModule.swift
index 91c2b055c..19f69a132 100644
--- a/ChatSecureTests/OTROMEMOTestModule.swift
+++ b/ChatSecureTests/OTROMEMOTestModule.swift
@@ -9,9 +9,15 @@
 import XCTest
 import XMPPFramework
 
+protocol OTROMEMOTestModuleProtocol: class {
+    func username() -> String
+    func receiveKeyData(keyData: [OMEMOKeyData], iv: NSData, fromJID: XMPPJID, senderDeviceId: gl_uint32_t, payload: NSData?, elementId: String?)
+    func bundle() -> OMEMOBundle
+}
+
 class OTROMEMOTestModule: OMEMOModule {
 
-    var otherUser:TestUser!
+    weak var otherUser:OTROMEMOTestModuleProtocol?
     var thisUser:TestUser!
     
     override var xmppStream:XMPPStream {
@@ -25,11 +31,11 @@ class OTROMEMOTestModule: OMEMOModule {
     /** Manually called after all the otherUser and thisUser are setup */
     override func xmppStreamDidAuthenticate(sender: XMPPStream!) {
         
-        let bundle = otherUser.signalOMEMOCoordinator.fetchMyBundle()
+        let bundle = otherUser?.bundle()
         XCTAssertNotNil(bundle)
         let device = bundle?.deviceId
         XCTAssertNotNil(device)
-        let otherJID = XMPPJID.jidWithString(otherUser.account.username)
+        let otherJID = XMPPJID.jidWithString(otherUser?.username())
         let ourJID = XMPPJID.jidWithString(thisUser.account.username)
         //After authentication fake receiving devices from other buddy
         self.omemoStorage.storeDeviceIds([NSNumber(unsignedInt: device!)], forJID: otherJID)
@@ -39,12 +45,12 @@ class OTROMEMOTestModule: OMEMOModule {
     /** When fetching a bundle for device and jid we return that device given to us in the other user struct*/
     override func fetchBundleForDeviceId(deviceId: gl_uint32_t, jid: XMPPJID, elementId: String?) {
         dispatch_async(self.moduleQueue) { 
-            if self.otherUser.signalOMEMOCoordinator.fetchMyBundle()?.deviceId == deviceId {
+            if self.otherUser?.bundle().deviceId == deviceId {
                 let multicastDelegate = self.valueForKey("multicastDelegate")!
                 //Empty responses so not nil and have correct elementID.
                 let response = XMPPIQ(type: "get", to: nil, elementID: elementId, child: nil)
                 let outgoing = XMPPIQ(type: "get", to: nil, elementID: elementId, child: nil)
-                multicastDelegate.omemo!(self, fetchedBundle: self.otherUser.signalOMEMOCoordinator.fetchMyBundle()!, fromJID: jid, responseIq: response, outgoingIq: outgoing)
+                multicastDelegate.omemo!(self, fetchedBundle: self.otherUser!.bundle(), fromJID: jid, responseIq: response, outgoingIq: outgoing)
                 
             }
         }
@@ -52,8 +58,9 @@ class OTROMEMOTestModule: OMEMOModule {
     
     /** When we send key data we automtically route that data to the other user to decrypto*/
     override func sendKeyData(keyData: [OMEMOKeyData], iv: NSData, toJID: XMPPJID, payload: NSData?, elementId: String?) {
-        let dummyMessage = XMPPMessage(type: "chat", elementID: "1234")
-        self.otherUser.signalOMEMOCoordinator.omemo(self, receivedKeyData: keyData, iv: iv, senderDeviceId: (thisUser.signalOMEMOCoordinator.fetchMyBundle()?.deviceId)!, fromJID: XMPPJID.jidWithString(thisUser.account.username), payload: payload, message: dummyMessage)
+        
+        self.otherUser?.receiveKeyData(keyData, iv: iv, fromJID: XMPPJID.jidWithString(thisUser.account.username), senderDeviceId:(thisUser.signalOMEMOCoordinator.fetchMyBundle()?.deviceId)!, payload: payload, elementId: elementId)
+        //self.otherUser.signalOMEMOCoordinator.omemo(self, receivedKeyData: keyData, iv: iv, senderDeviceId: (thisUser.signalOMEMOCoordinator.fetchMyBundle()?.deviceId)!, fromJID: XMPPJID.jidWithString(thisUser.account.username), payload: payload, message: dummyMessage)
     }
     
     override func removeDeviceIds(deviceIds: [NSNumber], elementId: String?) {
@@ -64,3 +71,18 @@ class OTROMEMOTestModule: OMEMOModule {
         }
     }
 }
+
+extension OTROMEMOTestModule: OTROMEMOTestModuleProtocol {
+    func username() -> String {
+        return self.thisUser.account.username
+    }
+    
+    func receiveKeyData(keyData: [OMEMOKeyData], iv: NSData, fromJID: XMPPJID, senderDeviceId:gl_uint32_t, payload: NSData?, elementId: String?) {
+        let dummyMessage = XMPPMessage(type: "chat", elementID: "1234")
+        self.thisUser.signalOMEMOCoordinator.omemo(self, receivedKeyData: keyData, iv: iv, senderDeviceId: senderDeviceId, fromJID: fromJID, payload: payload, message: dummyMessage)
+    }
+    
+    func bundle() -> OMEMOBundle {
+        return self.thisUser.signalOMEMOCoordinator.fetchMyBundle()!
+    }
+}