#40: Add basic AuthenticationService and missing UI tests.

* Add MockAuthenticationService and ServerSelectionUITests. * Add tests covering the Authentication flow.
2026-11-03 06:02:10 +00:00 · 2022-07-04 10:00:27 +01:00
parent 0a199ee611
commit 26d8affaa8
29 changed files with 641 additions and 203 deletions
@@ -89,23 +89,15 @@ class AppCoordinator: AuthenticationCoordinatorDelegate, Coordinator {
    
    // MARK: - AuthenticationCoordinatorDelegate
    
-    func authenticationCoordinatorDidStartLoading(_ authenticationCoordinator: AuthenticationCoordinator) {
-        stateMachine.processEvent(.attemptedSignIn)
-    }
-    
    func authenticationCoordinator(_ authenticationCoordinator: AuthenticationCoordinator, didLoginWithSession userSession: UserSessionProtocol) {
        self.userSession = userSession
        remove(childCoordinator: authenticationCoordinator)
        stateMachine.processEvent(.succeededSigningIn)
    }
    
-    func authenticationCoordinator(_ authenticationCoordinator: AuthenticationCoordinator, didFailWithError error: AuthenticationCoordinatorError) {
-        stateMachine.processEvent(.failedSigningIn)
-    }
-    
    // MARK: - Private
    
-    // swiftlint:disable cyclomatic_complexity function_body_length
+    // swiftlint:disable cyclomatic_complexity
    private func setupStateMachine() {
        stateMachine.addTransitionHandler { [weak self] context in
            guard let self = self else { return }
@@ -113,13 +105,7 @@ class AppCoordinator: AuthenticationCoordinatorDelegate, Coordinator {
            switch (context.fromState, context.event, context.toState) {
            case (.initial, .startWithAuthentication, .signedOut):
                self.startAuthentication()
-            case (.signedOut, .attemptedSignIn, .signingIn):
-                self.showLoadingIndicator()
-            case (.signingIn, .failedSigningIn, .signedOut):
-                self.hideLoadingIndicator()
-                self.showLoginErrorToast()
-            case (.signingIn, .succeededSigningIn, .homeScreen):
-                self.hideLoadingIndicator()
+            case (.signedOut, .succeededSigningIn, .homeScreen):
                self.presentHomeScreen()
                
            case (.initial, .startWithExistingSession, .restoringSession):
@@ -179,7 +165,8 @@ class AppCoordinator: AuthenticationCoordinatorDelegate, Coordinator {
    }
    
    private func startAuthentication() {
-        let coordinator = AuthenticationCoordinator(userSessionStore: userSessionStore,
+        let authenticationService = AuthenticationService(userSessionStore: userSessionStore)
+        let coordinator = AuthenticationCoordinator(authenticationService: authenticationService,
                                                    navigationRouter: navigationRouter)
        coordinator.delegate = self
        
@@ -16,8 +16,6 @@ class AppCoordinatorStateMachine {
        case initial
        /// Showing the login screen
        case signedOut
-        /// Processing sign in request
-        case signingIn
        /// Opening an existing session.
        case restoringSession
        /// Showing the home screen
@@ -41,12 +39,8 @@ class AppCoordinatorStateMachine {
    enum Event: EventType {
        /// Start the `AppCoordinator` by showing authentication.
        case startWithAuthentication
-        /// A sign in request has been started
-        case attemptedSignIn
        /// Signing in succeeded
        case succeededSigningIn
-        /// Signing in failed
-        case failedSigningIn
        
        /// Start the `AppCoordinator` by restoring an existing account.
        case startWithExistingSession
@@ -84,9 +78,7 @@ class AppCoordinatorStateMachine {
    init() {
        stateMachine = StateMachine(state: .initial) { machine in
            machine.addRoutes(event: .startWithAuthentication, transitions: [ .initial => .signedOut ])
-            machine.addRoutes(event: .attemptedSignIn, transitions: [ .signedOut => .signingIn ])
-            machine.addRoutes(event: .succeededSigningIn, transitions: [ .signingIn => .homeScreen ])
-            machine.addRoutes(event: .failedSigningIn, transitions: [ .signingIn => .signedOut ])
+            machine.addRoutes(event: .succeededSigningIn, transitions: [ .signedOut => .homeScreen ])
            
            machine.addRoutes(event: .startWithExistingSession, transitions: [ .initial => .restoringSession ])
            machine.addRoutes(event: .succeededRestoringSession, transitions: [ .restoringSession => .homeScreen ])
@@ -31,4 +31,12 @@ extension String {
        
        return string
    }
+    
+    /// Whether or not the string is a Matrix user ID.
+    var isMatrixUserID: Bool {
+        let range = NSRange(location: 0, length: count)
+        
+        let detector = try? NSRegularExpression(pattern: MatrixEntityRegex.userId.rawValue, options: .caseInsensitive)
+        return detector?.numberOfMatches(in: self, range: range) ?? 0 == 1
+    }
 }
@@ -6,28 +6,18 @@
 //  Copyright © 2022 Element. All rights reserved.
 //

-import Foundation
+import UIKit
 import MatrixRustSDK

-enum AuthenticationCoordinatorError: Error {
-    case failedLoggingIn
-}
-
@MainActor
 protocol AuthenticationCoordinatorDelegate: AnyObject {
-    
-    func authenticationCoordinatorDidStartLoading(_ authenticationCoordinator: AuthenticationCoordinator)
-    
    func authenticationCoordinator(_ authenticationCoordinator: AuthenticationCoordinator,
                                   didLoginWithSession userSession: UserSessionProtocol)
-    
-    func authenticationCoordinator(_ authenticationCoordinator: AuthenticationCoordinator,
-                                   didFailWithError error: AuthenticationCoordinatorError)
 }

-class AuthenticationCoordinator: Coordinator {
+class AuthenticationCoordinator: Coordinator, Presentable {
    
-    private let userSessionStore: UserSessionStoreProtocol
+    private let authenticationService: AuthenticationServiceProtocol
    private let navigationRouter: NavigationRouter
    
    private(set) var clientProxy: ClientProxyProtocol?
@@ -35,9 +25,9 @@ class AuthenticationCoordinator: Coordinator {
    
    weak var delegate: AuthenticationCoordinatorDelegate?
    
-    init(userSessionStore: UserSessionStoreProtocol,
+    init(authenticationService: AuthenticationServiceProtocol,
         navigationRouter: NavigationRouter) {
-        self.userSessionStore = userSessionStore
+        self.authenticationService = authenticationService
        self.navigationRouter = navigationRouter
    }
    
@@ -45,6 +35,10 @@ class AuthenticationCoordinator: Coordinator {
        showSplashScreen()
    }
    
+    func toPresentable() -> UIViewController {
+        navigationRouter.toPresentable()
+    }
+    
    // MARK: - Private
    
    private func showSplashScreen() {
@@ -67,28 +61,18 @@ class AuthenticationCoordinator: Coordinator {
    }
    
    private func showLoginScreen() {
-        let homeserver = LoginHomeserver(address: BuildSettings.defaultHomeserverURLString)
-        let parameters = LoginCoordinatorParameters(navigationRouter: navigationRouter, homeserver: homeserver)
+        let parameters = LoginCoordinatorParameters(authenticationService: authenticationService,
+                                                    navigationRouter: navigationRouter)
        let coordinator = LoginCoordinator(parameters: parameters)
        
        coordinator.callback = { [weak self, weak coordinator] action in
-            guard let self = self, let coordinator = coordinator else {
-                return
-            }
+            guard let self = self, let coordinator = coordinator else { return }
            
            switch action {
-            case .login(let username, let password):
-                Task {
-                    switch await self.login(username: username, password: password) {
-                    case .success(let userSession):
-                        self.delegate?.authenticationCoordinator(self, didLoginWithSession: userSession)
-                        self.remove(childCoordinator: coordinator)
-                        self.navigationRouter.dismissModule()
-                    case .failure(let error):
-                        self.delegate?.authenticationCoordinator(self, didFailWithError: error)
-                        MXLog.error("Failed logging in user with error: \(error)")
-                    }
-                }
+            case .signedIn(let userSession):
+                self.delegate?.authenticationCoordinator(self, didLoginWithSession: userSession)
+                self.remove(childCoordinator: coordinator)
+                self.navigationRouter.dismissModule()
            case .continueWithOIDC:
                break
            }
@@ -101,38 +85,4 @@ class AuthenticationCoordinator: Coordinator {
            self?.remove(childCoordinator: coordinator)
        }
    }
-    
-    private func login(username: String, password: String) async -> Result<UserSession, AuthenticationCoordinatorError> {
-        Benchmark.startTrackingForIdentifier("Login", message: "Started new login")
-        
-        delegate?.authenticationCoordinatorDidStartLoading(self)
-        
-        let basePath = userSessionStore.baseDirectoryPath(for: username)
-        let builder = ClientBuilder()
-            .basePath(path: basePath)
-            .username(username: username)
-        
-        let loginTask: Task<Client, Error> = Task.detached {
-            let client = try builder.build()
-            try client.login(username: username, password: password)
-            return client
-        }
-        
-        switch await loginTask.result {
-        case .success(let client):
-            return await userSession(for: client)
-        case .failure(let error):
-            MXLog.error("Failed logging in with error: \(error)")
-            return .failure(.failedLoggingIn)
-        }
-    }
-    
-    private func userSession(for client: Client) async -> Result<UserSession, AuthenticationCoordinatorError> {
-        switch await userSessionStore.userSession(for: client) {
-        case .success(let clientProxy):
-            return .success(clientProxy)
-        case .failure:
-            return .failure(.failedLoggingIn)
-        }
-    }
 }
@@ -18,26 +18,17 @@ import SwiftUI
 import MatrixRustSDK

 struct LoginCoordinatorParameters {
+    /// The service used to authenticate the user.
+    let authenticationService: AuthenticationServiceProtocol
+    /// The navigation router used to present the server selection screen.
    let navigationRouter: NavigationRouterType
-    /// The homeserver to be shown initially.
-    let homeserver: LoginHomeserver
 }

-enum LoginCoordinatorAction: CustomStringConvertible {
-    /// Login with the associated username and password.
-    case login(username: String, password: String)
+enum LoginCoordinatorAction {
+    /// Login was successful.
+    case signedIn(UserSessionProtocol)
    /// Continue using OIDC.
    case continueWithOIDC
-    
-    /// A string representation of the action, ignoring any associated values that could leak PII.
-    var description: String {
-        switch self {
-        case .login:
-            return "login"
-        case .continueWithOIDC:
-            return "continueWithOIDC"
-        }
-    }
 }

 final class LoginCoordinator: Coordinator, Presentable {
@@ -56,6 +47,7 @@ final class LoginCoordinator: Coordinator, Presentable {
        }
    }
    
+    private var authenticationService: AuthenticationServiceProtocol { parameters.authenticationService }
    private var navigationRouter: NavigationRouterType { parameters.navigationRouter }
    private var indicatorPresenter: UserIndicatorTypePresenterProtocol
    private var activityIndicator: UserIndicator?
@@ -71,7 +63,7 @@ final class LoginCoordinator: Coordinator, Presentable {
    init(parameters: LoginCoordinatorParameters) {
        self.parameters = parameters
        
-        let viewModel = LoginViewModel(homeserver: parameters.homeserver)
+        let viewModel = LoginViewModel(homeserver: parameters.authenticationService.homeserver)
        loginViewModel = viewModel
        
        let view = LoginScreen(context: viewModel.context)
@@ -135,56 +127,68 @@ final class LoginCoordinator: Coordinator, Presentable {
    }
    
    /// Processes an error to either update the flow or display it to the user.
-    private func handleError(_ error: Error) {
-        loginViewModel.displayError(.alert(error.localizedDescription))
+    private func handleError(_ error: AuthenticationServiceError) {
+        switch error {
+        case .invalidCredentials:
+            loginViewModel.displayError(.alert(ElementL10n.authInvalidLoginParam))
+        case .accountDeactivated:
+            loginViewModel.displayError(.alert(ElementL10n.authInvalidLoginDeactivatedAccount))
+        default:
+            loginViewModel.displayError(.alert(ElementL10n.unknownError))
+        }
    }
    
    /// Requests the authentication coordinator to log in using the specified credentials.
    private func login(username: String, password: String) {
-        var username = loginViewModel.context.username
+        startLoading(isInteractionBlocking: true)
        
-        if !isMXID(username: username) {
-            let homeserver = loginViewModel.context.viewState.homeserver
-            username = "@\(username):\(homeserver.address)"
+        Task {
+            switch await authenticationService.login(username: username, password: password) {
+            case .success(let userSession):
+                callback?(.signedIn(userSession))
+                stopLoading()
+            case .failure(let error):
+                stopLoading()
+                handleError(error)
+            }
        }
-        
-        callback?(.login(username: username, password: password))
    }
    
    /// Parses the specified username and looks up the homeserver when a Matrix ID is entered.
    private func parseUsername(_ username: String) {
-        guard isMXID(username: username) else { return }
+        guard username.isMatrixUserID else { return }
        
-        let domain = String(username.split(separator: ":")[1])
+        let homeserverDomain = String(username.split(separator: ":")[1])
        
-        let homeserver = LoginHomeserver(address: domain)
-        updateViewModel(homeserver: homeserver)
-        indicateSuccess()
-    }
-    
-    /// Checks whether the specified username is a Matrix ID or not.
-    private func isMXID(username: String) -> Bool {
-        let range = NSRange(location: 0, length: username.count)
+        startLoading(isInteractionBlocking: false)
        
-        let detector = try? NSRegularExpression(pattern: MatrixEntityRegex.userId.rawValue, options: .caseInsensitive)
-        return detector?.numberOfMatches(in: username, range: range) ?? 0 > 0
+        Task {
+            switch await authenticationService.startLogin(for: homeserverDomain) {
+            case .success:
+                updateViewModel()
+                stopLoading()
+            case .failure(let error):
+                stopLoading()
+                handleError(error)
+            }
+        }
    }
    
    /// Updates the view model with a different homeserver.
-    private func updateViewModel(homeserver: LoginHomeserver) {
-        loginViewModel.update(homeserver: homeserver)
+    private func updateViewModel() {
+        loginViewModel.update(homeserver: authenticationService.homeserver)
        indicateSuccess()
    }
    
    /// Presents the server selection screen as a modal.
    private func presentServerSelectionScreen() {
        MXLog.debug("[LoginCoordinator] presentServerSelectionScreen")
-        let parameters = ServerSelectionCoordinatorParameters(homeserver: loginViewModel.context.viewState.homeserver,
+        let parameters = ServerSelectionCoordinatorParameters(authenticationService: authenticationService,
                                                              hasModalPresentation: true)
        let coordinator = ServerSelectionCoordinator(parameters: parameters)
-        coordinator.callback = { [weak self, weak coordinator] result in
+        coordinator.callback = { [weak self, weak coordinator] action in
            guard let self = self, let coordinator = coordinator else { return }
-            self.serverSelectionCoordinator(coordinator, didCompleteWith: result)
+            self.serverSelectionCoordinator(coordinator, didCompleteWith: action)
        }
        
        coordinator.start()
@@ -198,10 +202,10 @@ final class LoginCoordinator: Coordinator, Presentable {
    
    /// Handles the result from the server selection modal, dismissing it after updating the view.
    private func serverSelectionCoordinator(_ coordinator: ServerSelectionCoordinator,
-                                            didCompleteWith result: ServerSelectionCoordinatorResult) {
+                                            didCompleteWith action: ServerSelectionCoordinatorAction) {
        navigationRouter.dismissModule(animated: true) { [weak self] in
-            if case let .selected(homeserver) = result {
-                self?.updateViewModel(homeserver: homeserver)
+            if action == .updated {
+                self?.updateViewModel()
            }

            self?.remove(childCoordinator: coordinator)
@@ -58,6 +58,7 @@ struct LoginServerInfoSection: View {
                        .padding(.vertical, 2)
                }
                .buttonStyle(.elementGhost())
+                .accessibilityIdentifier("editServerButton")
            }
        }
    }
@@ -17,14 +17,14 @@
 import SwiftUI

 struct ServerSelectionCoordinatorParameters {
-    /// The homeserver to be shown initially.
-    let homeserver: LoginHomeserver
+    /// The service used to authenticate the user.
+    let authenticationService: AuthenticationServiceProtocol
    /// Whether the screen is presented modally or within a navigation stack.
    let hasModalPresentation: Bool
 }

-enum ServerSelectionCoordinatorResult {
-    case selected(LoginHomeserver)
+enum ServerSelectionCoordinatorAction {
+    case updated
    case dismiss
 }

@@ -38,6 +38,7 @@ final class ServerSelectionCoordinator: Coordinator, Presentable {
    private let serverSelectionHostingController: UIViewController
    private var serverSelectionViewModel: ServerSelectionViewModelProtocol
    
+    private var authenticationService: AuthenticationServiceProtocol { parameters.authenticationService }
    private var indicatorPresenter: UserIndicatorTypePresenterProtocol
    private var loadingIndicator: UserIndicator?
    
@@ -45,14 +46,14 @@ final class ServerSelectionCoordinator: Coordinator, Presentable {

    // Must be used only internally
    var childCoordinators: [Coordinator] = []
-    var callback: (@MainActor (ServerSelectionCoordinatorResult) -> Void)?
+    var callback: (@MainActor (ServerSelectionCoordinatorAction) -> Void)?
    
    // MARK: - Setup
    
    init(parameters: ServerSelectionCoordinatorParameters) {
        self.parameters = parameters
        
-        let viewModel = ServerSelectionViewModel(homeserverAddress: parameters.homeserver.address,
+        let viewModel = ServerSelectionViewModel(homeserverAddress: parameters.authenticationService.homeserver.address,
                                                 hasModalPresentation: parameters.hasModalPresentation)
        let view = ServerSelectionScreen(context: viewModel.context)
        serverSelectionViewModel = viewModel
@@ -66,11 +67,11 @@ final class ServerSelectionCoordinator: Coordinator, Presentable {
    func start() {
        MXLog.debug("[ServerSelectionCoordinator] did start.")
        
-        serverSelectionViewModel.callback = { [weak self] result in
+        serverSelectionViewModel.callback = { [weak self] action in
            guard let self = self else { return }
-            MXLog.debug("[ServerSelectionCoordinator] ServerSelectionViewModel did complete with result: \(result).")
+            MXLog.debug("[ServerSelectionCoordinator] ServerSelectionViewModel did callback with action: \(action).")
            
-            switch result {
+            switch action {
            case .confirm(let homeserverAddress):
                self.useHomeserver(homeserverAddress)
            case .dismiss:
@@ -102,9 +103,25 @@ final class ServerSelectionCoordinator: Coordinator, Presentable {
    private func useHomeserver(_ homeserverAddress: String) {
        startLoading()
        
-        let homeserverAddress = LoginHomeserver.sanitized(homeserverAddress)
-        
-        stopLoading()
-        callback?(.selected(LoginHomeserver(address: homeserverAddress)))
+        Task {
+            switch await authenticationService.startLogin(for: homeserverAddress) {
+            case .success:
+                callback?(.updated)
+                stopLoading()
+            case .failure(let error):
+                stopLoading()
+                handleError(error)
+            }
+        }
+    }
+    
+    /// Processes an error to either update the flow or display it to the user.
+    private func handleError(_ error: AuthenticationServiceError) {
+        switch error {
+        case .invalidServer:
+            serverSelectionViewModel.displayError(.footerMessage(ElementL10n.loginErrorHomeserverNotFound))
+        default:
+            serverSelectionViewModel.displayError(.footerMessage(ElementL10n.unknownError))
+        }
    }
 }
@@ -18,7 +18,7 @@ import Foundation

 // MARK: View model

-enum ServerSelectionViewModelResult {
+enum ServerSelectionViewModelAction {
    /// The user would like to use the homeserver at the given address.
    case confirm(homeserverAddress: String)
    /// Dismiss the view without using the entered address.
@@ -26,7 +26,7 @@ class ServerSelectionViewModel: ServerSelectionViewModelType, ServerSelectionVie

    // MARK: Public

-    var callback: (@MainActor (ServerSelectionViewModelResult) -> Void)?
+    var callback: (@MainActor (ServerSelectionViewModelAction) -> Void)?

    // MARK: - Setup

@@ -19,7 +19,7 @@ import Foundation
@MainActor
 protocol ServerSelectionViewModelProtocol {
    
-    var callback: (@MainActor (ServerSelectionViewModelResult) -> Void)? { get set }
+    var callback: (@MainActor (ServerSelectionViewModelAction) -> Void)? { get set }
    var context: ServerSelectionViewModelType.Context { get }
    
    /// Displays an error to the user.
@@ -0,0 +1,82 @@
+//
+//  AuthenticationService.swift
+//  ElementX
+//
+//  Created by Doug on 29/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Foundation
+import MatrixRustSDK
+
+class AuthenticationService: AuthenticationServiceProtocol {
+    
+    // MARK: - Properties
+    
+    // MARK: Private
+    
+    private(set) var homeserver: LoginHomeserver = LoginHomeserver(address: BuildSettings.defaultHomeserverURLString)
+    private let userSessionStore: UserSessionStoreProtocol
+    
+    // MARK: - Setup
+    
+    init(userSessionStore: UserSessionStoreProtocol) {
+        self.userSessionStore = userSessionStore
+    }
+    
+    // MARK: - Public
+    
+    func startLogin(for homeserverAddress: String) async -> Result<Void, AuthenticationServiceError> {
+        homeserver = LoginHomeserver(address: homeserverAddress)
+        return .success(())
+    }
+    
+    func login(username: String, password: String) async -> Result<UserSessionProtocol, AuthenticationServiceError> {
+        Benchmark.startTrackingForIdentifier("Login", message: "Started new login")
+        
+        // Workaround whilst the SDK requires a full MXID.
+        let username = username.isMatrixUserID ? username : "@\(username):\(homeserver.address)"
+        
+        let basePath = userSessionStore.baseDirectoryPath(for: username)
+        let builder = ClientBuilder()
+            .basePath(path: basePath)
+            .username(username: username)
+        
+        let loginTask: Task<Client, Error> = Task.detached {
+            let client = try builder.build()
+            try client.login(username: username, password: password)
+            return client
+        }
+        
+        switch await loginTask.result {
+        case .success(let client):
+            Benchmark.endTrackingForIdentifier("Login", message: "Finished login")
+            return await userSession(for: client)
+        case .failure(let error):
+            Benchmark.endTrackingForIdentifier("Login", message: "Login failed")
+            
+            MXLog.error("Failed logging in with error: \(error)")
+            guard let error = error as? ClientError else { return .failure(.failedLoggingIn) }
+            
+            switch error.code {
+            case .forbidden:
+                return .failure(.invalidCredentials)
+            case .userDeactivated:
+                return .failure(.accountDeactivated)
+            default:
+                return .failure(.failedLoggingIn)
+            }
+        }
+    }
+    
+    // MARK: - Private
+    
+    private func userSession(for client: Client) async -> Result<UserSessionProtocol, AuthenticationServiceError> {
+        switch await userSessionStore.userSession(for: client) {
+        case .success(let clientProxy):
+            return .success(clientProxy)
+        case .failure:
+            return .failure(.failedLoggingIn)
+        }
+    }
+}
@@ -0,0 +1,26 @@
+//
+//  AuthenticationServiceProtocol.swift
+//  ElementX
+//
+//  Created by Doug on 29/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Foundation
+
+enum AuthenticationServiceError: Error {
+    case invalidServer
+    case invalidCredentials
+    case accountDeactivated
+    case failedLoggingIn
+}
+
+@MainActor
+protocol AuthenticationServiceProtocol {
+    var homeserver: LoginHomeserver { get }
+    
+    /// Sets up the service for login on the specified homeserver address.
+    func startLogin(for homeserverAddress: String) async -> Result<Void, AuthenticationServiceError>
+    /// Performs a password login using the current homeserver.
+    func login(username: String, password: String) async -> Result<UserSessionProtocol, AuthenticationServiceError>
+}
@@ -0,0 +1,45 @@
+//
+//  MockAuthenticationService.swift
+//  ElementX
+//
+//  Created by Doug on 29/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Foundation
+
+class MockAuthenticationService: AuthenticationServiceProtocol {
+    let validCredentials = (username: "alice", password: "12345678")
+    private(set) var homeserver: LoginHomeserver = .mockMatrixDotOrg
+    
+    func startLogin(for homeserverAddress: String) async -> Result<Void, AuthenticationServiceError> {
+        // Map the address to the mock homeservers
+        if LoginHomeserver.mockMatrixDotOrg.address.contains(homeserverAddress) {
+            homeserver = .mockMatrixDotOrg
+            return .success(())
+        } else if LoginHomeserver.mockOIDC.address.contains(homeserverAddress) {
+            homeserver = .mockOIDC
+            return .success(())
+        } else if LoginHomeserver.mockBasicServer.address.contains(homeserverAddress) {
+            homeserver = .mockBasicServer
+            return .success(())
+        } else if LoginHomeserver.mockUnsupported.address.contains(homeserverAddress) {
+            homeserver = .mockUnsupported
+            return .success(())
+        } else {
+            // Otherwise fail with an invalid server.
+            return .failure(.invalidServer)
+        }
+    }
+    
+    func login(username: String, password: String) async -> Result<UserSessionProtocol, AuthenticationServiceError> {
+        // Login only succeeds if the username and password match the valid credentials property
+        guard username == validCredentials.username, password == validCredentials.password else {
+            return .failure(.invalidCredentials)
+        }
+        
+        let userSession = MockUserSession(clientProxy: MockClientProxy(userIdentifier: username),
+                                          mediaProvider: MockMediaProvider())
+        return .success(userSession)
+    }
+}
@@ -0,0 +1,30 @@
+//
+//  ClientError.swift
+//  ElementX
+//
+//  Created by Doug on 30/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Foundation
+import MatrixRustSDK
+
+enum MatrixErrorCode: String, CaseIterable {
+    case unknown = "M_UNKNOWN"
+    case userDeactivated = "M_USER_DEACTIVATED"
+    case forbidden = "M_FORBIDDEN"
+}
+
+extension ClientError {
+    var code: MatrixErrorCode {
+        guard case let .Generic(message) = self else { return .unknown }
+        
+        for code in MatrixErrorCode.allCases {
+            if message.contains(code.rawValue) {
+                return code
+            }
+        }
+        
+        return .unknown
+    }
+}
@@ -19,6 +19,7 @@ enum ClientProxyError: Error {
    case failedRetrievingAvatarURL
    case failedRetrievingDisplayName
    case failedRetrievingSessionVerificationController
+    case failedLoadingMedia
 }

 protocol ClientProxyProtocol {
@@ -0,0 +1,39 @@
+//
+//  MockClientProxy.swift
+//  ElementX
+//
+//  Created by Doug on 29/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Combine
+import MatrixRustSDK
+
+struct MockClientProxy: ClientProxyProtocol {
+    
+    let callbacks = PassthroughSubject<ClientProxyCallback, Never>()
+    
+    let userIdentifier: String
+    
+    let rooms = [RoomProxy]()
+    
+    func loadUserDisplayName() async -> Result<String, ClientProxyError> {
+        .failure(.failedRetrievingDisplayName)
+    }
+    
+    func loadUserAvatarURLString() async -> Result<String, ClientProxyError> {
+        .failure(.failedRetrievingAvatarURL)
+    }
+    
+    func mediaSourceForURLString(_ urlString: String) -> MatrixRustSDK.MediaSource {
+        MatrixRustSDK.mediaSourceFromUrl(url: urlString)
+    }
+    
+    func loadMediaContentForSource(_ source: MatrixRustSDK.MediaSource) throws -> Data {
+        throw ClientProxyError.failedLoadingMedia
+    }
+    
+    func sessionVerificationControllerProxy() async -> Result<SessionVerificationControllerProxyProtocol, ClientProxyError> {
+        .failure(.failedRetrievingSessionVerificationController)
+    }
+}
@@ -0,0 +1,17 @@
+//
+//  MockUserSession.swift
+//  ElementX
+//
+//  Created by Doug on 29/06/2022.
+//  Copyright © 2022 Element. All rights reserved.
+//
+
+import Combine
+
+struct MockUserSession: UserSessionProtocol {
+    let callbacks = PassthroughSubject<UserSessionCallback, Never>()
+    let sessionVerificationController: SessionVerificationControllerProxyProtocol? = nil
+    
+    let clientProxy: ClientProxyProtocol
+    let mediaProvider: MediaProviderProtocol
+}
@@ -98,8 +98,6 @@ class UserSessionStore: UserSessionStoreProtocol {
    }
    
    private func setupProxyForClient(_ client: Client) async -> Result<ClientProxyProtocol, UserSessionStoreError> {
-        Benchmark.endTrackingForIdentifier("Login", message: "Finished login")
-        
        do {
            let accessToken = try client.restoreToken()
            let userId = try client.userId()
@@ -10,8 +10,9 @@ import Foundation

 enum UITestScreenIdentifier: String {
    case login
-    case loginOIDC
-    case loginUnsupported
+    case serverSelection
+    case serverSelectionNonModal
+    case authenticationFlow
    case simpleRegular
    case simpleUpgrade
    case settings
@@ -11,28 +11,37 @@ import SwiftUI

 class UITestsAppCoordinator: Coordinator {
    private let window: UIWindow
-    private let mainNavigationController: UINavigationController
+    private let mainNavigationController: ElementNavigationController
+    private let navigationRouter: NavigationRouter
+    private var hostingController: UIViewController?
    
    var childCoordinators: [Coordinator] = []
    
    init() {
-        mainNavigationController = UINavigationController()
+        mainNavigationController = ElementNavigationController()
+        navigationRouter = NavigationRouter(navigationController: mainNavigationController)
+        
        window = UIWindow(frame: UIScreen.main.bounds)
        window.rootViewController = mainNavigationController
        window.tintColor = .element.accent
        
+        UIView.setAnimationsEnabled(false)
+        
        let screens = mockScreens()
+        
        let rootView = UITestsRootView(mockScreens: screens) { id in
            guard let screen = screens.first(where: { $0.id == id }) else {
                fatalError()
            }
            
            screen.coordinator.start()
-            
-            self.mainNavigationController.pushViewController(screen.coordinator.toPresentable(), animated: true)
+            self.navigationRouter.setRootModule(screen.coordinator)
        }
        
-        mainNavigationController.setViewControllers([UIHostingController(rootView: rootView)], animated: false)
+        let hostingController = UIHostingController(rootView: rootView)
+        self.hostingController = hostingController
+        
+        mainNavigationController.setViewControllers([hostingController], animated: false)
    }
    
    func start() {
@@ -40,34 +49,34 @@ class UITestsAppCoordinator: Coordinator {
    }
    
    private func mockScreens() -> [MockScreen] {
-        UITestScreenIdentifier.allCases.map { MockScreen(id: $0) }
+        UITestScreenIdentifier.allCases.map { MockScreen(id: $0, navigationRouter: navigationRouter) }
    }
 }

@MainActor
 class MockScreen: Identifiable {
    let id: UITestScreenIdentifier
+    let navigationRouter: NavigationRouter
    lazy var coordinator: Coordinator & Presentable = {
        switch id {
        case .login:
-            let router = NavigationRouter(navigationController: ElementNavigationController())
-            return LoginCoordinator(parameters: .init(navigationRouter: router,
-                                                      homeserver: .mockMatrixDotOrg))
-        case .loginOIDC:
-            let router = NavigationRouter(navigationController: ElementNavigationController())
-            return LoginCoordinator(parameters: .init(navigationRouter: router,
-                                                      homeserver: .mockOIDC))
-        case .loginUnsupported:
-            let router = NavigationRouter(navigationController: ElementNavigationController())
-            return LoginCoordinator(parameters: .init(navigationRouter: router,
-                                                      homeserver: .mockUnsupported))
+            return LoginCoordinator(parameters: .init(authenticationService: MockAuthenticationService(),
+                                                      navigationRouter: navigationRouter))
+        case .serverSelection:
+            return ServerSelectionCoordinator(parameters: .init(authenticationService: MockAuthenticationService(),
+                                                                hasModalPresentation: true))
+        case .serverSelectionNonModal:
+            return ServerSelectionCoordinator(parameters: .init(authenticationService: MockAuthenticationService(),
+                                                                hasModalPresentation: false))
+        case .authenticationFlow:
+            return AuthenticationCoordinator(authenticationService: MockAuthenticationService(),
+                                             navigationRouter: navigationRouter)
        case .simpleRegular:
            return TemplateCoordinator(parameters: .init(promptType: .regular))
        case .simpleUpgrade:
            return TemplateCoordinator(parameters: .init(promptType: .upgrade))
        case .settings:
-            let router = NavigationRouter(navigationController: ElementNavigationController())
-            return SettingsCoordinator(parameters: .init(navigationRouter: router,
+            return SettingsCoordinator(parameters: .init(navigationRouter: navigationRouter,
                                                         bugReportService: MockBugReportService()))
        case .bugReport:
            return BugReportCoordinator(parameters: .init(bugReportService: MockBugReportService(),
@@ -95,7 +104,8 @@ class MockScreen: Identifiable {
        }
    }()
    
-    init(id: UITestScreenIdentifier) {
+    init(id: UITestScreenIdentifier, navigationRouter: NavigationRouter) {
        self.id = id
+        self.navigationRouter = navigationRouter
    }
 }