Add ServiceConnectionManager extensions

4 files changed, 38 insertions, 24 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountRepository.kt
index 4d79bfd356..6a58739aec 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountRepository.kt
@@ -23,9 +23,6 @@ class AccountRepository(
     private val serviceConnectionManager: ServiceConnectionManager,
     dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) {
-    private val dataSource
-        get() = serviceConnectionManager.connectionState.value.readyContainer()?.accountDataSource
-
     private val _cachedCreatedAccount = MutableStateFlow<String?>(null)
     val cachedCreatedAccount = _cachedCreatedAccount.asStateFlow()
 
@@ -78,28 +75,28 @@ class AccountRepository(
         )
 
     fun createAccount() {
-        dataSource?.createAccount()
+        serviceConnectionManager.accountDataSource()?.createAccount()
     }
 
     fun login(accountToken: String) {
-        dataSource?.login(accountToken)
+        serviceConnectionManager.accountDataSource()?.login(accountToken)
     }
 
     fun logout() {
         clearCreatedAccountCache()
-        dataSource?.logout()
+        serviceConnectionManager.accountDataSource()?.logout()
     }
 
     fun fetchAccountExpiry() {
-        dataSource?.fetchAccountExpiry()
+        serviceConnectionManager.accountDataSource()?.fetchAccountExpiry()
     }
 
     fun fetchAccountHistory() {
-        dataSource?.fetchAccountHistory()
+        serviceConnectionManager.accountDataSource()?.fetchAccountHistory()
     }
 
     fun clearAccountHistory() {
-        dataSource?.clearAccountHistory()
+        serviceConnectionManager.accountDataSource()?.clearAccountHistory()
     }
 
     private fun clearCreatedAccountCache() {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
index 08290ef7d2..790c2404f2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
@@ -61,20 +61,16 @@ class DeviceRepository(
         .stateIn(CoroutineScope(Dispatchers.IO), SharingStarted.WhileSubscribed(), emptyList())
 
     fun refreshDeviceState() {
-        container()?.deviceDataSource?.refreshDevice()
-    }
-
-    private fun container(): ServiceConnectionContainer? {
-        return serviceConnectionManager.connectionState.value.readyContainer()
+        serviceConnectionManager.deviceDataSource()?.refreshDevice()
     }
 
     fun removeDevice(accountToken: String, deviceId: String) {
         cachedDeviceList.value = emptyList()
-        container()?.deviceDataSource?.removeDevice(accountToken, deviceId)
+        serviceConnectionManager.deviceDataSource()?.removeDevice(accountToken, deviceId)
     }
 
     fun refreshDeviceList(accountToken: String) {
-        container()?.deviceDataSource?.refreshDeviceList(accountToken)
+        serviceConnectionManager.deviceDataSource()?.refreshDeviceList(accountToken)
     }
 
     suspend fun getDeviceList(accountToken: String): DeviceListEvent {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionManagerExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionManagerExtensions.kt
new file mode 100644
index 0000000000..392b841101
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionManagerExtensions.kt
@@ -0,0 +1,25 @@
+package net.mullvad.mullvadvpn.ui.serviceconnection
+
+fun ServiceConnectionManager.accountDataSource() =
+    this.connectionState.value.readyContainer()?.accountDataSource
+
+fun ServiceConnectionManager.appVersionInfoCache() =
+    this.connectionState.value.readyContainer()?.appVersionInfoCache
+
+fun ServiceConnectionManager.authTokenCache() =
+    this.connectionState.value.readyContainer()?.authTokenCache
+
+fun ServiceConnectionManager.connectionProxy() =
+    this.connectionState.value.readyContainer()?.connectionProxy
+
+fun ServiceConnectionManager.deviceDataSource() =
+    this.connectionState.value.readyContainer()?.deviceDataSource
+
+fun ServiceConnectionManager.customDns() =
+    this.connectionState.value.readyContainer()?.customDns
+
+fun ServiceConnectionManager.relayListListener() =
+    this.connectionState.value.readyContainer()?.relayListListener
+
+fun ServiceConnectionManager.settingsListener() =
+    this.connectionState.value.readyContainer()?.settingsListener
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
index 4272a5de35..bf8db5b273 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
@@ -11,8 +11,8 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountRepository
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
+import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
 import net.mullvad.talpid.util.callbackFlowFromSubscription
 
 // TODO: Refactor ConnectionProxy to be easily injectable rather than injecting
@@ -44,15 +44,11 @@ class DeviceRevokedViewModel(
         )
 
     fun onGoToLoginClicked() {
-        serviceContainer()?.let { container ->
-            if (container.connectionProxy.state.isSecured()) {
-                container.connectionProxy.disconnect()
+        serviceConnectionManager.connectionProxy()?.let { proxy ->
+            if (proxy.state.isSecured()) {
+                proxy.disconnect()
             }
             accountRepository.logout()
         }
     }
-
-    private fun serviceContainer(): ServiceConnectionContainer? {
-        return serviceConnectionManager.connectionState.value.readyContainer()
-    }
 }