From ada8d82a3c35f7a02af1119ddbbcdbb2a7efb397 Mon Sep 17 00:00:00 2001
From: Dariusz Rybicki <dariusz@elixxir.io>
Date: Mon, 22 Aug 2022 13:02:53 +0100
Subject: [PATCH] Update MessengerContext

---
 .../Messenger/Functors/MessengerConnect.swift |   6 +-
 .../Functors/MessengerIsConnected.swift       |   2 +-
 .../Functors/MessengerIsLoaded.swift          |   2 +-
 .../Functors/MessengerIsLoggedIn.swift        |   2 +-
 .../Functors/MessengerIsRegistered.swift      |   2 +-
 .../Messenger/Functors/MessengerLoad.swift    |   4 +-
 .../Messenger/Functors/MessengerLogIn.swift   |   8 +-
 .../Functors/MessengerRegister.swift          |   8 +-
 .../Messenger/MessengerContext.swift          |  49 ++++++--
 .../Messenger/MessengerEnvironment.swift      |  42 ++++---
 .../Functors/MessengerConnectTests.swift      |  39 +++---
 .../Functors/MessengerIsConnectedTests.swift  |   8 +-
 .../Functors/MessengerIsLoadedTests.swift     |   8 +-
 .../Functors/MessengerIsLoggedInTests.swift   |   8 +-
 .../Functors/MessengerIsRegisteredTests.swift |  26 ++--
 .../Functors/MessengerLoadTests.swift         |   7 +-
 .../Functors/MessengerLogInTests.swift        | 101 ++++++++++-----
 .../Functors/MessengerRegisterTests.swift     | 118 +++++++++++-------
 18 files changed, 276 insertions(+), 164 deletions(-)

diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerConnect.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerConnect.swift
index fa24b8eb..97e399ab 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerConnect.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerConnect.swift
@@ -16,14 +16,14 @@ public struct MessengerConnect {
 extension MessengerConnect {
   public static func live(_ env: MessengerEnvironment) -> MessengerConnect {
     MessengerConnect {
-      guard let cMix = env.ctx.cMix else {
+      guard let cMix = env.ctx.getCMix() else {
         throw Error.notLoaded
       }
-      env.ctx.e2e = try env.login(
+      env.ctx.setE2E(try env.login(
         cMixId: cMix.getId(),
         identity: try cMix.makeLegacyReceptionIdentity(),
         e2eParamsJSON: env.getE2EParams()
-      )
+      ))
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsConnected.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsConnected.swift
index 6a790dd5..81017429 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsConnected.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsConnected.swift
@@ -12,7 +12,7 @@ public struct MessengerIsConnected {
 extension MessengerIsConnected {
   public static func live(_ env: MessengerEnvironment) -> MessengerIsConnected {
     MessengerIsConnected {
-      env.ctx.e2e != nil
+      env.ctx.getE2E() != nil
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoaded.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoaded.swift
index 88750885..8931d2b4 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoaded.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoaded.swift
@@ -12,7 +12,7 @@ public struct MessengerIsLoaded {
 extension MessengerIsLoaded {
   public static func live(_ env: MessengerEnvironment) -> MessengerIsLoaded {
     MessengerIsLoaded {
-      env.ctx.cMix != nil
+      env.ctx.getCMix() != nil
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoggedIn.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoggedIn.swift
index 7f886bef..04aaa338 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoggedIn.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsLoggedIn.swift
@@ -12,7 +12,7 @@ public struct MessengerIsLoggedIn {
 extension MessengerIsLoggedIn {
   public static func live(_ env: MessengerEnvironment) -> MessengerIsLoggedIn {
     MessengerIsLoggedIn {
-      env.ctx.ud != nil
+      env.ctx.getUD() != nil
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsRegistered.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsRegistered.swift
index c91db5ac..cbf02c8a 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerIsRegistered.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerIsRegistered.swift
@@ -16,7 +16,7 @@ public struct MessengerIsRegistered {
 extension MessengerIsRegistered {
   public static func live(_ env: MessengerEnvironment) -> MessengerIsRegistered {
     MessengerIsRegistered {
-      guard let e2e = env.ctx.e2e else {
+      guard let e2e = env.ctx.getE2E() else {
         throw Error.notConnected
       }
       return try env.isRegisteredWithUD(e2eId: e2e.getId())
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerLoad.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerLoad.swift
index a8a6be31..c57d7a5a 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerLoad.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerLoad.swift
@@ -12,11 +12,11 @@ public struct MessengerLoad {
 extension MessengerLoad {
   public static func live(_ env: MessengerEnvironment) -> MessengerLoad {
     MessengerLoad {
-      env.ctx.cMix = try env.loadCMix(
+      env.ctx.setCMix(try env.loadCMix(
         storageDir: env.storageDir,
         password: try env.passwordStorage.load(),
         cMixParamsJSON: env.getCMixParams()
-      )
+      ))
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerLogIn.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerLogIn.swift
index 70df78ac..31323e9f 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerLogIn.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerLogIn.swift
@@ -17,16 +17,16 @@ public struct MessengerLogIn {
 extension MessengerLogIn {
   public static func live(_ env: MessengerEnvironment) -> MessengerLogIn {
     MessengerLogIn {
-      guard let cMix = env.ctx.cMix else {
+      guard let cMix = env.ctx.getCMix() else {
         throw Error.notLoaded
       }
-      guard let e2e = env.ctx.e2e else {
+      guard let e2e = env.ctx.getE2E() else {
         throw Error.notConnected
       }
       if cMix.networkFollowerStatus() != .running {
         try cMix.startNetworkFollower(timeoutMS: 30_000)
       }
-      env.ctx.ud = try env.newOrLoadUd(
+      env.ctx.setUD(try env.newOrLoadUd(
         params: .init(
           e2eId: e2e.getId(),
           username: nil,
@@ -38,7 +38,7 @@ extension MessengerLogIn {
         follower: .init {
           cMix.networkFollowerStatus().rawValue
         }
-      )
+      ))
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/Functors/MessengerRegister.swift b/Sources/XXMessengerClient/Messenger/Functors/MessengerRegister.swift
index 75fce252..0fdddac8 100644
--- a/Sources/XXMessengerClient/Messenger/Functors/MessengerRegister.swift
+++ b/Sources/XXMessengerClient/Messenger/Functors/MessengerRegister.swift
@@ -19,16 +19,16 @@ public struct MessengerRegister {
 extension MessengerRegister {
   public static func live(_ env: MessengerEnvironment) -> MessengerRegister {
     MessengerRegister { username in
-      guard let cMix = env.ctx.cMix else {
+      guard let cMix = env.ctx.getCMix() else {
         throw Error.notLoaded
       }
-      guard let e2e = env.ctx.e2e else {
+      guard let e2e = env.ctx.getE2E() else {
         throw Error.notConnected
       }
       if cMix.networkFollowerStatus() != .running {
         try cMix.startNetworkFollower(timeoutMS: 30_000)
       }
-      env.ctx.ud = try env.newOrLoadUd(
+      env.ctx.setUD(try env.newOrLoadUd(
         params: .init(
           e2eId: e2e.getId(),
           username: username,
@@ -40,7 +40,7 @@ extension MessengerRegister {
         follower: .init {
           cMix.networkFollowerStatus().rawValue
         }
-      )
+      ))
     }
   }
 }
diff --git a/Sources/XXMessengerClient/Messenger/MessengerContext.swift b/Sources/XXMessengerClient/Messenger/MessengerContext.swift
index ba67ed2b..c98523f2 100644
--- a/Sources/XXMessengerClient/Messenger/MessengerContext.swift
+++ b/Sources/XXMessengerClient/Messenger/MessengerContext.swift
@@ -1,17 +1,42 @@
 import XXClient
+import XCTestDynamicOverlay
 
-public class MessengerContext {
-  public init(
-    cMix: CMix? = nil,
-    e2e: E2E? = nil,
-    ud: UserDiscovery? = nil
-  ) {
-    self.cMix = cMix
-    self.e2e = e2e
-    self.ud = ud
+public struct MessengerContext {
+  public var getCMix: () -> CMix?
+  public var setCMix: (CMix?) -> Void
+  public var getE2E: () -> E2E?
+  public var setE2E: (E2E?) -> Void
+  public var getUD: () -> UserDiscovery?
+  public var setUD: (UserDiscovery?) -> Void
+}
+
+extension MessengerContext {
+  public static func live() -> MessengerContext {
+    class Container {
+      var cMix: CMix?
+      var e2e: E2E?
+      var ud: UserDiscovery?
+    }
+    let container = Container()
+
+    return MessengerContext(
+      getCMix: { container.cMix },
+      setCMix: { container.cMix = $0 },
+      getE2E: { container.e2e },
+      setE2E: { container.e2e = $0 },
+      getUD: { container.ud },
+      setUD: { container.ud = $0 }
+    )
   }
+}
 
-  public var cMix: CMix?
-  public var e2e: E2E?
-  public var ud: UserDiscovery?
+extension MessengerContext {
+  public static let unimplemented = MessengerContext(
+    getCMix: XCTUnimplemented("\(Self.self).getCMix"),
+    setCMix: XCTUnimplemented("\(Self.self).setCMix"),
+    getE2E: XCTUnimplemented("\(Self.self).getE2E"),
+    setE2E: XCTUnimplemented("\(Self.self).setE2E"),
+    getUD: XCTUnimplemented("\(Self.self).getUD"),
+    setUD: XCTUnimplemented("\(Self.self).setUD")
+  )
 }
diff --git a/Sources/XXMessengerClient/Messenger/MessengerEnvironment.swift b/Sources/XXMessengerClient/Messenger/MessengerEnvironment.swift
index 29391b20..b3009ae0 100644
--- a/Sources/XXMessengerClient/Messenger/MessengerEnvironment.swift
+++ b/Sources/XXMessengerClient/Messenger/MessengerEnvironment.swift
@@ -29,30 +29,32 @@ extension MessengerEnvironment {
     .appendingPathComponent("xx.network.client")
     .path
 
-  public static let live = MessengerEnvironment(
-    ctx: .init(),
-    downloadNDF: .live,
-    fileManager: .live(),
-    generateSecret: .live,
-    getCMixParams: .liveDefault,
-    getE2EParams: .liveDefault,
-    isRegisteredWithUD: .live,
-    loadCMix: .live,
-    login: .live,
-    ndfEnvironment: .mainnet,
-    newCMix: .live,
-    newOrLoadUd: .live,
-    passwordStorage: .keychain,
-    storageDir: MessengerEnvironment.defaultStorageDir,
-    udAddress: nil,
-    udCert: nil,
-    udContact: nil
-  )
+  public static func live() -> MessengerEnvironment {
+    MessengerEnvironment(
+      ctx: .live(),
+      downloadNDF: .live,
+      fileManager: .live(),
+      generateSecret: .live,
+      getCMixParams: .liveDefault,
+      getE2EParams: .liveDefault,
+      isRegisteredWithUD: .live,
+      loadCMix: .live,
+      login: .live,
+      ndfEnvironment: .mainnet,
+      newCMix: .live,
+      newOrLoadUd: .live,
+      passwordStorage: .keychain,
+      storageDir: MessengerEnvironment.defaultStorageDir,
+      udAddress: nil,
+      udCert: nil,
+      udContact: nil
+    )
+  }
 }
 
 extension MessengerEnvironment {
   public static let unimplemented = MessengerEnvironment(
-    ctx: .init(),
+    ctx: .unimplemented,
     downloadNDF: .unimplemented,
     fileManager: .unimplemented,
     generateSecret: .unimplemented,
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerConnectTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerConnectTests.swift
index c44b4907..c93e403e 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerConnectTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerConnectTests.swift
@@ -12,16 +12,22 @@ final class MessengerConnectTests: XCTestCase {
       var identity: ReceptionIdentity
       var e2eParamsJSON: Data
     }
+
     var didLogIn: [DidLogIn] = []
+    var didSetE2E: [E2E?] = []
 
     let cMixId = 1234
     let receptionId = ReceptionIdentity.stub
     let e2eParams = "e2e-params".data(using: .utf8)!
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.getId.run = { cMixId }
-    env.ctx.cMix!.makeLegacyReceptionIdentity.run = { receptionId }
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.getId.run = { cMixId }
+      cMix.makeLegacyReceptionIdentity.run = { receptionId }
+      return cMix
+    }
+    env.ctx.setE2E = { didSetE2E.append($0) }
     env.getE2EParams.run = { e2eParams }
     env.login.run = { ephemeral, cMixId, authCallbacks, identity, e2eParamsJSON in
       didLogIn.append(.init(
@@ -46,13 +52,12 @@ final class MessengerConnectTests: XCTestCase {
         e2eParamsJSON: e2eParams
       )
     ])
-
-    XCTAssertNotNil(env.ctx.e2e)
+    XCTAssertEqual(didSetE2E.compactMap { $0 }.count, 1)
   }
 
   func testConnectWithoutCMix() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { nil }
     let connect: MessengerConnect = .live(env)
 
     XCTAssertThrowsError(try connect()) { error in
@@ -67,10 +72,13 @@ final class MessengerConnectTests: XCTestCase {
     struct Error: Swift.Error, Equatable {}
     let error = Error()
 
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.getId.run = { 1234 }
-    env.ctx.cMix!.makeLegacyReceptionIdentity.run = { throw error }
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.getId.run = { 1234 }
+      cMix.makeLegacyReceptionIdentity.run = { throw error }
+      return cMix
+    }
     let connect: MessengerConnect = .live(env)
 
     XCTAssertThrowsError(try connect()) { err in
@@ -83,9 +91,12 @@ final class MessengerConnectTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.getId.run = { 1234 }
-    env.ctx.cMix!.makeLegacyReceptionIdentity.run = { .stub }
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.getId.run = { 1234 }
+      cMix.makeLegacyReceptionIdentity.run = { .stub }
+      return cMix
+    }
     env.getE2EParams.run = { "e2e-params".data(using: .utf8)! }
     env.login.run = { _, _, _, _, _ in throw error }
     let connect: MessengerConnect = .live(env)
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsConnectedTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsConnectedTests.swift
index e207e8a5..b361170c 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsConnectedTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsConnectedTests.swift
@@ -3,16 +3,16 @@ import XCTest
 
 final class MessengerIsConnectedTests: XCTestCase {
   func testWithE2E() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = .unimplemented
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getE2E = { .unimplemented }
     let isConnected: MessengerIsConnected = .live(env)
 
     XCTAssertTrue(isConnected())
   }
 
   func testWithoutE2E() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getE2E = { nil }
     let isConnected: MessengerIsConnected = .live(env)
 
     XCTAssertFalse(isConnected())
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoadedTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoadedTests.swift
index c83dc523..a88a1148 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoadedTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoadedTests.swift
@@ -3,16 +3,16 @@ import XCTest
 
 final class MessengerIsLoadedTests: XCTestCase {
   func testWithCMix() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { .unimplemented }
     let isLoaded: MessengerIsLoaded = .live(env)
 
     XCTAssertTrue(isLoaded())
   }
 
   func testWithoutCMix() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { nil }
     let isLoaded: MessengerIsLoaded = .live(env)
 
     XCTAssertFalse(isLoaded())
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoggedInTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoggedInTests.swift
index 6b5b7903..0bbd8867 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoggedInTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsLoggedInTests.swift
@@ -3,16 +3,16 @@ import XCTest
 
 final class MessengerIsLoggedInTests: XCTestCase {
   func testWithUD() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.ud = .unimplemented
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getUD = { .unimplemented }
     let isLoggedIn: MessengerIsLoggedIn = .live(env)
 
     XCTAssertTrue(isLoggedIn())
   }
 
   func testWithoutUD() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.ud = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getUD = { nil }
     let isLoggedIn: MessengerIsLoggedIn = .live(env)
 
     XCTAssertFalse(isLoggedIn())
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsRegisteredTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsRegisteredTests.swift
index 250fdad8..b11bfbc6 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsRegisteredTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerIsRegisteredTests.swift
@@ -1,4 +1,5 @@
 import XCTest
+import XXClient
 @testable import XXMessengerClient
 
 final class MessengerIsRegisteredTests: XCTestCase {
@@ -6,8 +7,11 @@ final class MessengerIsRegisteredTests: XCTestCase {
     var didIsRegisteredWithUD: [Int] = []
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      return e2e
+    }
     env.isRegisteredWithUD.run = { e2eId in
       didIsRegisteredWithUD.append(e2eId)
       return true
@@ -20,8 +24,11 @@ final class MessengerIsRegisteredTests: XCTestCase {
 
   func testNotRegistered() throws {
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      return e2e
+    }
     env.isRegisteredWithUD.run = { _ in false }
     let isRegistered: MessengerIsRegistered = .live(env)
 
@@ -29,8 +36,8 @@ final class MessengerIsRegisteredTests: XCTestCase {
   }
 
   func testWithoutE2E() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getE2E = { nil }
     let isRegistered: MessengerIsRegistered = .live(env)
 
     XCTAssertThrowsError(try isRegistered()) { err in
@@ -46,8 +53,11 @@ final class MessengerIsRegisteredTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      return e2e
+    }
     env.isRegisteredWithUD.run = { _ in throw error }
     let isRegistered: MessengerIsRegistered = .live(env)
 
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLoadTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLoadTests.swift
index 1d8ae5f2..eecdca12 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLoadTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLoadTests.swift
@@ -11,13 +11,14 @@ final class MessengerLoadTests: XCTestCase {
       var cMixParamsJSON: Data
     }
     var didLoadCMix: [DidLoadCMix] = []
+    var didSetCMix: [CMix?] = []
 
     let storageDir = "test-storage-dir"
     let password = "password".data(using: .utf8)!
     let cMixParams = "cmix-params".data(using: .utf8)!
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
+    env.ctx.setCMix = { didSetCMix.append($0) }
     env.storageDir = storageDir
     env.passwordStorage.load = { password }
     env.getCMixParams.run = { cMixParams }
@@ -40,12 +41,11 @@ final class MessengerLoadTests: XCTestCase {
         cMixParamsJSON: cMixParams
       )
     ])
-    XCTAssertNotNil(env.ctx.cMix)
+    XCTAssertEqual(didSetCMix.compactMap{ $0 }.count, 1)
   }
 
   func testMissingPassword() {
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
     env.storageDir = "storage-dir"
     env.passwordStorage.load = { throw PasswordStorage.MissingPasswordError() }
     let load: MessengerLoad = .live(env)
@@ -63,7 +63,6 @@ final class MessengerLoadTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
     env.storageDir = "storage-dir"
     env.passwordStorage.load = { "password".data(using: .utf8)! }
     env.getCMixParams.run = { "cmix-params".data(using: .utf8)! }
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLogInTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLogInTests.swift
index b11f8c8f..791c28ff 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLogInTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerLogInTests.swift
@@ -8,6 +8,7 @@ final class MessengerLogInTests: XCTestCase {
     var didStartNetworkFollower: [Int] = []
     var didNewOrLoadUDWithParams: [NewOrLoadUd.Params] = []
     var didNewOrLoadUDWithFollower: [UdNetworkStatus] = []
+    var didSetUD: [UserDiscovery?] = []
 
     let e2eId = 1234
     let networkFollowerStatus: NetworkFollowerStatus = .stopped
@@ -16,19 +17,26 @@ final class MessengerLogInTests: XCTestCase {
     let udAddressFromNDF = "ndf-ud-address"
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { networkFollowerStatus }
-    env.ctx.cMix!.startNetworkFollower.run = { timeout in
-      didStartNetworkFollower.append(timeout)
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { networkFollowerStatus }
+      cMix.startNetworkFollower.run = { timeout in
+        didStartNetworkFollower.append(timeout)
+      }
+      return cMix
     }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { e2eId }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { e2eId }
+      e2e.getUdCertFromNdf.run = { udCertFromNDF }
+      e2e.getUdContactFromNdf.run = { udContactFromNDF }
+      e2e.getUdAddressFromNdf.run = { udAddressFromNDF }
+      return e2e
+    }
+    env.ctx.setUD = { didSetUD.append($0) }
     env.udCert = nil
     env.udContact = nil
     env.udAddress = nil
-    env.ctx.e2e!.getUdCertFromNdf.run = { udCertFromNDF }
-    env.ctx.e2e!.getUdContactFromNdf.run = { udContactFromNDF }
-    env.ctx.e2e!.getUdAddressFromNdf.run = { udAddressFromNDF }
     env.newOrLoadUd.run = { params, follower in
       didNewOrLoadUDWithParams.append(params)
       didNewOrLoadUDWithFollower.append(follower)
@@ -51,21 +59,30 @@ final class MessengerLogInTests: XCTestCase {
       didNewOrLoadUDWithFollower.first?.handle(),
       networkFollowerStatus.rawValue
     )
-    XCTAssertNotNil(env.ctx.ud)
+    XCTAssertEqual(didSetUD.compactMap { $0 }.count, 1)
   }
 
   func testLoginWithAlternativeUD() throws {
     var didNewOrLoadUDWithParams: [NewOrLoadUd.Params] = []
+    var didSetUD: [UserDiscovery?] = []
+
     let e2eId = 1234
     let altUdCert = "alt-ud-cert".data(using: .utf8)!
     let altUdContact = "alt-ud-contact".data(using: .utf8)!
     let altUdAddress = "alt-ud-address"
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { e2eId }
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { e2eId }
+      return e2e
+    }
+    env.ctx.setUD = { didSetUD.append($0) }
     env.udCert = altUdCert
     env.udContact = altUdContact
     env.udAddress = altUdAddress
@@ -84,12 +101,12 @@ final class MessengerLogInTests: XCTestCase {
       contactFile: altUdContact,
       address: altUdAddress
     )])
-    XCTAssertNotNil(env.ctx.ud)
+    XCTAssertEqual(didSetUD.compactMap { $0 }.count, 1)
   }
 
   func testLoginWithoutCMix() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { nil }
     let logIn: MessengerLogIn = .live(env)
 
     XCTAssertThrowsError(try logIn()) { error in
@@ -101,9 +118,9 @@ final class MessengerLogInTests: XCTestCase {
   }
 
   func testLoginWithoutE2E() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.e2e = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { .unimplemented }
+    env.ctx.getE2E = { nil }
     let logIn: MessengerLogIn = .live(env)
 
     XCTAssertThrowsError(try logIn()) { error in
@@ -118,10 +135,14 @@ final class MessengerLogInTests: XCTestCase {
     struct Error: Swift.Error, Equatable {}
     let error = Error()
 
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .stopped }
-    env.ctx.cMix!.startNetworkFollower.run = { _ in throw error }
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .stopped }
+      cMix.startNetworkFollower.run = { _ in throw error }
+      return cMix
+    }
+    env.ctx.getE2E = { .unimplemented }
     let logIn: MessengerLogIn = .live(env)
 
     XCTAssertThrowsError(try logIn()) { err in
@@ -134,14 +155,20 @@ final class MessengerLogInTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      e2e.getUdCertFromNdf.run = { "ndf-ud-cert".data(using: .utf8)! }
+      e2e.getUdContactFromNdf.run = { throw error }
+      return e2e
+    }
     env.udCert = nil
     env.udContact = nil
-    env.ctx.e2e!.getUdCertFromNdf.run = { "ndf-ud-cert".data(using: .utf8)! }
-    env.ctx.e2e!.getUdContactFromNdf.run = { throw error }
     let logIn: MessengerLogIn = .live(env)
 
     XCTAssertThrowsError(try logIn()) { err in
@@ -154,10 +181,16 @@ final class MessengerLogInTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      return e2e
+    }
     env.udCert = "ud-cert".data(using: .utf8)!
     env.udContact = "ud-contact".data(using: .utf8)!
     env.udAddress = "ud-address"
diff --git a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerRegisterTests.swift b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerRegisterTests.swift
index 861df63e..8feb8a7c 100644
--- a/Tests/XXMessengerClientTests/Messenger/Functors/MessengerRegisterTests.swift
+++ b/Tests/XXMessengerClientTests/Messenger/Functors/MessengerRegisterTests.swift
@@ -8,6 +8,7 @@ final class MessengerRegisterTests: XCTestCase {
     var didStartNetworkFollower: [Int] = []
     var didNewOrLoadUDWithParams: [NewOrLoadUd.Params] = []
     var didNewOrLoadUDWithFollower: [UdNetworkStatus] = []
+    var didSetUD: [UserDiscovery?] = []
 
     let e2eId = 1234
     let networkFollowerStatus: NetworkFollowerStatus = .stopped
@@ -18,22 +19,29 @@ final class MessengerRegisterTests: XCTestCase {
     let username = "new-user-name"
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { networkFollowerStatus }
-    env.ctx.cMix!.startNetworkFollower.run = { timeout in
-      didStartNetworkFollower.append(timeout)
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { networkFollowerStatus }
+      cMix.startNetworkFollower.run = { timeout in
+        didStartNetworkFollower.append(timeout)
+      }
+      cMix.getReceptionRegistrationValidationSignature.run = {
+        registrationSignature
+      }
+      return cMix
     }
-    env.ctx.cMix!.getReceptionRegistrationValidationSignature.run = {
-      registrationSignature
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { e2eId }
+      e2e.getUdCertFromNdf.run = { udCertFromNDF }
+      e2e.getUdContactFromNdf.run = { udContactFromNDF }
+      e2e.getUdAddressFromNdf.run = { udAddressFromNDF }
+      return e2e
     }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { e2eId }
+    env.ctx.setUD = { didSetUD.append($0) }
     env.udCert = nil
     env.udContact = nil
     env.udAddress = nil
-    env.ctx.e2e!.getUdCertFromNdf.run = { udCertFromNDF }
-    env.ctx.e2e!.getUdContactFromNdf.run = { udContactFromNDF }
-    env.ctx.e2e!.getUdAddressFromNdf.run = { udAddressFromNDF }
     env.newOrLoadUd.run = { params, follower in
       didNewOrLoadUDWithParams.append(params)
       didNewOrLoadUDWithFollower.append(follower)
@@ -56,11 +64,13 @@ final class MessengerRegisterTests: XCTestCase {
       didNewOrLoadUDWithFollower.first?.handle(),
       networkFollowerStatus.rawValue
     )
-    XCTAssertNotNil(env.ctx.ud)
+    XCTAssertEqual(didSetUD.compactMap { $0 }.count, 1)
   }
 
   func testRegisterWithAlternativeUD() throws {
     var didNewOrLoadUDWithParams: [NewOrLoadUd.Params] = []
+    var didSetUD: [UserDiscovery?] = []
+
     let e2eId = 1234
     let registrationSignature = "registration-signature".data(using: .utf8)!
     let altUdCert = "alt-ud-cert".data(using: .utf8)!
@@ -69,13 +79,20 @@ final class MessengerRegisterTests: XCTestCase {
     let username = "new-user-name"
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.cMix!.getReceptionRegistrationValidationSignature.run = {
-      registrationSignature
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      cMix.getReceptionRegistrationValidationSignature.run = {
+        registrationSignature
+      }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { e2eId }
+      return e2e
     }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { e2eId }
+    env.ctx.setUD = { didSetUD.append($0) }
     env.udCert = altUdCert
     env.udContact = altUdContact
     env.udAddress = altUdAddress
@@ -94,12 +111,12 @@ final class MessengerRegisterTests: XCTestCase {
       contactFile: altUdContact,
       address: altUdAddress
     )])
-    XCTAssertNotNil(env.ctx.ud)
+    XCTAssertEqual(didSetUD.compactMap { $0 }.count, 1)
   }
 
   func testRegisterWithoutCMix() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { nil }
     let register: MessengerRegister = .live(env)
 
     XCTAssertThrowsError(try register(username: "new-user-name")) { error in
@@ -111,9 +128,9 @@ final class MessengerRegisterTests: XCTestCase {
   }
 
   func testRegisterWithoutE2E() {
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.e2e = nil
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = { .unimplemented }
+    env.ctx.getE2E = { nil }
     let register: MessengerRegister = .live(env)
 
     XCTAssertThrowsError(try register(username: "new-user-name")) { error in
@@ -128,11 +145,14 @@ final class MessengerRegisterTests: XCTestCase {
     struct Error: Swift.Error, Equatable {}
     let error = Error()
 
-    let env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .stopped }
-    env.ctx.cMix!.startNetworkFollower.run = { _ in throw error }
-    env.ctx.e2e = .unimplemented
+    var env: MessengerEnvironment = .unimplemented
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .stopped }
+      cMix.startNetworkFollower.run = { _ in throw error }
+      return cMix
+    }
+    env.ctx.getE2E = { .unimplemented }
     let register: MessengerRegister = .live(env)
 
     XCTAssertThrowsError(try register(username: "new-user-name")) { err in
@@ -145,17 +165,23 @@ final class MessengerRegisterTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.cMix!.getReceptionRegistrationValidationSignature.run = {
-      "registration-signature".data(using: .utf8)!
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      cMix.getReceptionRegistrationValidationSignature.run = {
+        "registration-signature".data(using: .utf8)!
+      }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      e2e.getUdCertFromNdf.run = { "ndf-ud-cert".data(using: .utf8)! }
+      e2e.getUdContactFromNdf.run = { throw error }
+      return e2e
     }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
     env.udCert = nil
     env.udContact = nil
-    env.ctx.e2e!.getUdCertFromNdf.run = { "ndf-ud-cert".data(using: .utf8)! }
-    env.ctx.e2e!.getUdContactFromNdf.run = { throw error }
     let register: MessengerRegister = .live(env)
 
     XCTAssertThrowsError(try register(username: "new-user-name")) { err in
@@ -168,13 +194,19 @@ final class MessengerRegisterTests: XCTestCase {
     let error = Error()
 
     var env: MessengerEnvironment = .unimplemented
-    env.ctx.cMix = .unimplemented
-    env.ctx.cMix!.networkFollowerStatus.run = { .running }
-    env.ctx.cMix!.getReceptionRegistrationValidationSignature.run = {
-      "registration-signature".data(using: .utf8)!
+    env.ctx.getCMix = {
+      var cMix: CMix = .unimplemented
+      cMix.networkFollowerStatus.run = { .running }
+      cMix.getReceptionRegistrationValidationSignature.run = {
+        "registration-signature".data(using: .utf8)!
+      }
+      return cMix
+    }
+    env.ctx.getE2E = {
+      var e2e: E2E = .unimplemented
+      e2e.getId.run = { 1234 }
+      return e2e
     }
-    env.ctx.e2e = .unimplemented
-    env.ctx.e2e!.getId.run = { 1234 }
     env.udCert = "ud-cert".data(using: .utf8)!
     env.udContact = "ud-contact".data(using: .utf8)!
     env.udAddress = "ud-address"
-- 
GitLab