Toggling Internet Connectivity in Android Instrumentation Test

UiInternetConnectivitySettingsPanel.kt

import android.annotation.TargetApi
import android.content.Context
import android.content.Intent
import android.provider.Settings
import androidx.test.platform.app.InstrumentationRegistry
import androidx.test.uiautomator.By
import androidx.test.uiautomator.StaleObjectException
import androidx.test.uiautomator.UiDevice
import androidx.test.uiautomator.UiObject2
import androidx.test.uiautomator.Until

private const val UI_DEVICE_ELEMENT_TIMEOUT_MS = 10000L

/**
 * A test utility that enables configuring a device's network from a test using the Settings
 * panel API introduced in Android Q.
 *
 * Usage:
 *
 * ```
 * val uiConnectivitySettingsPanel = UiInternetConnectivitySettingsPanel(context)
 *
 * uiConnectivitySettingsPanel.enableWifi(true)
 *
 * uiConnectivitySettingsPanel.enableAirplaneMode(true)
 * ```
 */
@TargetApi(29)
class UiInternetConnectivitySettingsPanel(private val context: Context) {
    private val uiDevice by lazy {
        UiDevice.getInstance(InstrumentationRegistry.getInstrumentation())
    }

    fun enableWifi(enable: Boolean) {
        val originalAppPackage = launchInternetSettingsPanel()

        wifiSwitch().apply {
            if ((enable && !isChecked) || (!enable && isChecked)) {
                click()
            }
        }

        closeInternetSettingsPanel(originalAppPackage)
    }

    fun enableMobileData(enable: Boolean) {
        val originalAppPackage = launchInternetSettingsPanel()

        mobileDataSwitch().apply {
            if ((enable && !isChecked) || (!enable && isChecked)) {
                click()
            }
        }

        closeInternetSettingsPanel(originalAppPackage)
    }

    fun enableAirplaneMode(enable: Boolean) {
        val originalAppPackage = launchInternetSettingsPanel()

        /**
         * Set the switches as specified
         */
        airplaneModeSwitch().apply {
            if ((enable && !isChecked) || (!enable && isChecked)) {
                click()
                return
            }
        }

        closeInternetSettingsPanel(originalAppPackage)
    }

    private fun launchInternetSettingsPanel(): String {
        val originalAppPackage = uiDevice.currentPackageName

        /**
         * Launch and await the internet connectivity settings panel
         */
        context.startActivity(Intent(Settings.Panel.ACTION_INTERNET_CONNECTIVITY).apply {
            addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
        })
        uiDevice.wait(
            Until.findObject(By.pkg("com.android.settings")),
            UI_DEVICE_ELEMENT_TIMEOUT_MS
        )

        return originalAppPackage
    }

    private fun closeInternetSettingsPanel(originalAppPackage: String) {
        val doneButton = uiDevice.wait(Until.findObject(By.textContains("Done")),
            UI_DEVICE_ELEMENT_TIMEOUT_MS)
        doneButton.click()
        uiDevice.wait(
            Until.findObject(By.pkg(originalAppPackage)),
            UI_DEVICE_ELEMENT_TIMEOUT_MS
        )
    }

    private fun wifiSwitch(): UiObject2 {
        return retryUiSelector {
            findInternetConnectivitySwitch("Wi‑Fi")
        }
    }

    private fun mobileDataSwitch(): UiObject2 {
        return retryUiSelector {
            findInternetConnectivitySwitch("Mobile data")
        }
    }

    private fun airplaneModeSwitch(): UiObject2 {
        return retryUiSelector {
            findInternetConnectivitySwitch("Airplane mode")
        }
    }

    private fun findInternetConnectivitySwitch(name: String): UiObject2 {
        return uiDevice.wait(
            Until.findObject(By.textContains(name)),
            UI_DEVICE_ELEMENT_TIMEOUT_MS)
            .parent.parent.findObject(By.clazz("android.widget.Switch"))
    }

    private fun retryUiSelector(uiSelector: () -> UiObject2): UiObject2 {
        var uiObject2: UiObject2?
        var retries = 0
        val maxRetries = 5
        do {
            Thread.sleep(500)
            uiObject2 = try {
                uiSelector()
            } catch (e: StaleObjectException) {
                null
            }
        } while (uiObject2 == null && retries++ < maxRetries)

        requireNotNull(uiObject2)
        return uiObject2
    }
}