2 files changed, 18 insertions, 4 deletions
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
index dac1426b2a..39c535c1d7 100644
--- a/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/ConnectFragment.kt
@@ -26,6 +26,7 @@ import net.mullvad.mullvadvpn.model.TunnelStateTransition
 
 class ConnectFragment : Fragment() {
     private lateinit var actionButton: ConnectActionButton
+    private lateinit var switchLocationButton: SwitchLocationButton
     private lateinit var headerBar: HeaderBar
     private lateinit var notificationBanner: NotificationBanner
     private lateinit var status: ConnectionStatus
@@ -64,10 +65,6 @@ class ConnectFragment : Fragment() {
             parentActivity.openSettings()
         }
 
-        view.findViewById<Button>(R.id.switch_location).setOnClickListener {
-            openSwitchLocationScreen()
-        }
-
         headerBar = HeaderBar(view, context!!)
         notificationBanner = NotificationBanner(view)
         status = ConnectionStatus(view, context!!)
@@ -80,6 +77,9 @@ class ConnectFragment : Fragment() {
             onDisconnect = { disconnect() }
         }
 
+        switchLocationButton = SwitchLocationButton(view)
+        switchLocationButton.onClick = { openSwitchLocationScreen() }
+
         attachListenerJob = attachListener()
 
         return view
diff --git a/android/src/main/kotlin/net/mullvad/mullvadvpn/SwitchLocationButton.kt b/android/src/main/kotlin/net/mullvad/mullvadvpn/SwitchLocationButton.kt
new file mode 100644
index 0000000000..72a33dc295
--- /dev/null
+++ b/android/src/main/kotlin/net/mullvad/mullvadvpn/SwitchLocationButton.kt
@@ -0,0 +1,14 @@
+package net.mullvad.mullvadvpn
+
+import android.view.View
+import android.widget.Button
+
+class SwitchLocationButton(val parentView: View) {
+    private val button: Button = parentView.findViewById(R.id.switch_location)
+
+    var onClick: (() -> Unit)? = null
+
+    init {
+        button.setOnClickListener { onClick?.invoke() }
+    }
+}