[SB Refactor] Fix handling of network transports in wifi repo.

Fixes:
1) Handle CELLULAR transport first, like the old pipeline.
2) Don't assume `NetworkCapabilities.transportInfo` is always of type
   `WifiInfo`.
3) Update `isWifiDefault` to match the old pipeline. (If the default
   network has the wifi transport, the wifi is default even if the
   `transportInfo` is not of type `WifiInfo`.)

Bug: 266628069
Test: atest WifiRepositoryImplTest
Test: manual: verify wifi icon still works
Change-Id: I47ab4eb01ee8a3f893acbf13a9a01bbd9fdee10d

5 files changed, 157 insertions, 13 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarIconController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarIconController.java
index 24ad55d67bb0..11863627218e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarIconController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarIconController.java
@@ -501,7 +501,7 @@ public interface StatusBarIconController {
         @VisibleForTesting
         protected StatusIconDisplayable addWifiIcon(int index, String slot, WifiIconState state) {
             if (mStatusBarPipelineFlags.useNewWifiIcon()) {
-                throw new IllegalStateException("Attempting to add a mobile icon while the new "
+                throw new IllegalStateException("Attempting to add a wifi icon while the new "
                         + "icons are enabled is not supported");
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLogger.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLogger.kt
index d3ff3573dae2..491f3a5513b0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLogger.kt
@@ -97,15 +97,20 @@ constructor(
         )
     }
 
-    fun logOnCapabilitiesChanged(network: Network, networkCapabilities: NetworkCapabilities) {
+    fun logOnCapabilitiesChanged(
+        network: Network,
+        networkCapabilities: NetworkCapabilities,
+        isDefaultNetworkCallback: Boolean,
+    ) {
         buffer.log(
             SB_LOGGING_TAG,
             LogLevel.INFO,
             {
+                bool1 = isDefaultNetworkCallback
                 int1 = network.getNetId()
                 str1 = networkCapabilities.toString()
             },
-            { "onCapabilitiesChanged: net=$int1 capabilities=$str1" }
+            { "onCapabilitiesChanged[default=$bool1]: net=$int1 capabilities=$str1" }
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImpl.kt
index d26499c18661..86690479f679 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImpl.kt
@@ -114,13 +114,17 @@ constructor(
                             network: Network,
                             networkCapabilities: NetworkCapabilities
                         ) {
+                            logger.logOnCapabilitiesChanged(
+                                network,
+                                networkCapabilities,
+                                isDefaultNetworkCallback = true,
+                            )
+
                             // This method will always be called immediately after the network
                             // becomes the default, in addition to any time the capabilities change
                             // while the network is the default.
-                            // If this network contains valid wifi info, then wifi is the default
-                            // network.
-                            val wifiInfo = networkCapabilitiesToWifiInfo(networkCapabilities)
-                            trySend(wifiInfo != null)
+                            // If this network is a wifi network, then wifi is the default network.
+                            trySend(isWifiNetwork(networkCapabilities))
                         }
 
                         override fun onLost(network: Network) {
@@ -152,7 +156,11 @@ constructor(
                             network: Network,
                             networkCapabilities: NetworkCapabilities
                         ) {
-                            logger.logOnCapabilitiesChanged(network, networkCapabilities)
+                            logger.logOnCapabilitiesChanged(
+                                network,
+                                networkCapabilities,
+                                isDefaultNetworkCallback = false,
+                            )
 
                             wifiNetworkChangeEvents.tryEmit(Unit)
 
@@ -253,16 +261,30 @@ constructor(
             networkCapabilities: NetworkCapabilities
         ): WifiInfo? {
             return when {
-                networkCapabilities.hasTransport(TRANSPORT_WIFI) ->
-                    networkCapabilities.transportInfo as WifiInfo?
                 networkCapabilities.hasTransport(TRANSPORT_CELLULAR) ->
                     // Sometimes, cellular networks can act as wifi networks (known as VCN --
                     // virtual carrier network). So, see if this cellular network has wifi info.
                     Utils.tryGetWifiInfoForVcn(networkCapabilities)
+                networkCapabilities.hasTransport(TRANSPORT_WIFI) ->
+                    if (networkCapabilities.transportInfo is WifiInfo) {
+                        networkCapabilities.transportInfo as WifiInfo
+                    } else {
+                        null
+                    }
                 else -> null
             }
         }
 
+        /** True if these capabilities represent a wifi network. */
+        private fun isWifiNetwork(networkCapabilities: NetworkCapabilities): Boolean {
+            return when {
+                networkCapabilities.hasTransport(TRANSPORT_WIFI) -> true
+                networkCapabilities.hasTransport(TRANSPORT_CELLULAR) ->
+                    Utils.tryGetWifiInfoForVcn(networkCapabilities) != null
+                else -> false
+            }
+        }
+
         private fun createWifiNetworkModel(
             wifiInfo: WifiInfo,
             network: Network,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLoggerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLoggerTest.kt
index b32058fca109..3dccbbf26575 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLoggerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ConnectivityPipelineLoggerTest.kt
@@ -45,7 +45,7 @@ class ConnectivityPipelineLoggerTest : SysuiTestCase() {
 
     @Test
     fun testLogNetworkCapsChange_bufferHasInfo() {
-        logger.logOnCapabilitiesChanged(NET_1, NET_1_CAPS)
+        logger.logOnCapabilitiesChanged(NET_1, NET_1_CAPS, isDefaultNetworkCallback = true)
 
         val stringWriter = StringWriter()
         buffer.dump(PrintWriter(stringWriter), tailLength = 0)
@@ -54,6 +54,7 @@ class ConnectivityPipelineLoggerTest : SysuiTestCase() {
         val expectedNetId = NET_1_ID.toString()
         val expectedCaps = NET_1_CAPS.toString()
 
+        assertThat(actualString).contains("true")
         assertThat(actualString).contains(expectedNetId)
         assertThat(actualString).contains(expectedCaps)
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImplTest.kt
index 87ce8faff5a5..7099f1f0af2d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/data/repository/prod/WifiRepositoryImplTest.kt
@@ -21,7 +21,10 @@ import android.net.Network
 import android.net.NetworkCapabilities
 import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
 import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.net.NetworkCapabilities.TRANSPORT_VPN
 import android.net.NetworkCapabilities.TRANSPORT_WIFI
+import android.net.TransportInfo
+import android.net.VpnTransportInfo
 import android.net.vcn.VcnTransportInfo
 import android.net.wifi.WifiInfo
 import android.net.wifi.WifiManager
@@ -243,6 +246,54 @@ class WifiRepositoryImplTest : SysuiTestCase() {
         job.cancel()
     }
 
+    /** Regression test for b/266628069. */
+    @Test
+    fun isWifiDefault_transportInfoIsNotWifi_andNoWifiTransport_false() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.isWifiDefault.launchIn(this)
+
+            val transportInfo = VpnTransportInfo(
+                /* type= */ 0,
+                /* sessionId= */ "sessionId",
+            )
+            val networkCapabilities = mock<NetworkCapabilities>().also {
+                whenever(it.hasTransport(TRANSPORT_VPN)).thenReturn(true)
+                whenever(it.hasTransport(TRANSPORT_WIFI)).thenReturn(false)
+                whenever(it.hasTransport(TRANSPORT_CELLULAR)).thenReturn(false)
+                whenever(it.transportInfo).thenReturn(transportInfo)
+            }
+
+            getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, networkCapabilities)
+
+            assertThat(underTest.isWifiDefault.value).isFalse()
+
+            job.cancel()
+        }
+
+    /** Regression test for b/266628069. */
+    @Test
+    fun isWifiDefault_transportInfoIsNotWifi_butHasWifiTransport_true() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.isWifiDefault.launchIn(this)
+
+            val transportInfo = VpnTransportInfo(
+                /* type= */ 0,
+                /* sessionId= */ "sessionId",
+            )
+            val networkCapabilities = mock<NetworkCapabilities>().also {
+                whenever(it.hasTransport(TRANSPORT_VPN)).thenReturn(true)
+                whenever(it.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
+                whenever(it.hasTransport(TRANSPORT_CELLULAR)).thenReturn(false)
+                whenever(it.transportInfo).thenReturn(transportInfo)
+            }
+
+            getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, networkCapabilities)
+
+            assertThat(underTest.isWifiDefault.value).isTrue()
+
+            job.cancel()
+        }
+
     @Test
     fun isWifiDefault_cellularVcnNetwork_isTrue() = runBlocking(IMMEDIATE) {
         val job = underTest.isWifiDefault.launchIn(this)
@@ -260,6 +311,24 @@ class WifiRepositoryImplTest : SysuiTestCase() {
     }
 
     @Test
+    fun wifiNetwork_cellularAndWifiTransports_usesCellular_isTrue() =
+        runBlocking(IMMEDIATE) {
+            val job = underTest.isWifiDefault.launchIn(this)
+
+            val capabilities = mock<NetworkCapabilities>().apply {
+                whenever(this.hasTransport(TRANSPORT_CELLULAR)).thenReturn(true)
+                whenever(this.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
+                whenever(this.transportInfo).thenReturn(VcnTransportInfo(PRIMARY_WIFI_INFO))
+            }
+
+            getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
+
+            assertThat(underTest.isWifiDefault.value).isTrue()
+
+            job.cancel()
+        }
+
+    @Test
     fun isWifiDefault_cellularNotVcnNetwork_isFalse() = runBlocking(IMMEDIATE) {
         val job = underTest.isWifiDefault.launchIn(this)
 
@@ -467,6 +536,28 @@ class WifiRepositoryImplTest : SysuiTestCase() {
         job.cancel()
     }
 
+    /** Regression test for b/266628069. */
+    @Test
+    fun wifiNetwork_transportInfoIsNotWifi_flowHasNoNetwork() =
+        runBlocking(IMMEDIATE) {
+            var latest: WifiNetworkModel? = null
+            val job = underTest
+                .wifiNetwork
+                .onEach { latest = it }
+                .launchIn(this)
+
+            val transportInfo = VpnTransportInfo(
+                /* type= */ 0,
+                /* sessionId= */ "sessionId",
+            )
+            getNetworkCallback()
+                .onCapabilitiesChanged(NETWORK, createWifiNetworkCapabilities(transportInfo))
+
+            assertThat(latest is WifiNetworkModel.Inactive).isTrue()
+
+            job.cancel()
+        }
+
     @Test
     fun wifiNetwork_cellularVcnNetworkAdded_flowHasNetwork() = runBlocking(IMMEDIATE) {
         var latest: WifiNetworkModel? = null
@@ -535,6 +626,31 @@ class WifiRepositoryImplTest : SysuiTestCase() {
     }
 
     @Test
+    fun wifiNetwork_cellularAndWifiTransports_usesCellular() =
+        runBlocking(IMMEDIATE) {
+            var latest: WifiNetworkModel? = null
+            val job = underTest
+                .wifiNetwork
+                .onEach { latest = it }
+                .launchIn(this)
+
+            val capabilities = mock<NetworkCapabilities>().apply {
+                whenever(this.hasTransport(TRANSPORT_CELLULAR)).thenReturn(true)
+                whenever(this.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
+                whenever(this.transportInfo).thenReturn(VcnTransportInfo(PRIMARY_WIFI_INFO))
+            }
+
+            getNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
+
+            assertThat(latest is WifiNetworkModel.Active).isTrue()
+            val latestActive = latest as WifiNetworkModel.Active
+            assertThat(latestActive.networkId).isEqualTo(NETWORK_ID)
+            assertThat(latestActive.ssid).isEqualTo(SSID)
+
+            job.cancel()
+        }
+
+    @Test
     fun wifiNetwork_newPrimaryWifiNetwork_flowHasNewNetwork() = runBlocking(IMMEDIATE) {
         var latest: WifiNetworkModel? = null
         val job = underTest
@@ -870,12 +986,12 @@ class WifiRepositoryImplTest : SysuiTestCase() {
     }
 
     private fun createWifiNetworkCapabilities(
-        wifiInfo: WifiInfo,
+        transportInfo: TransportInfo,
         isValidated: Boolean = true,
     ): NetworkCapabilities {
         return mock<NetworkCapabilities>().also {
             whenever(it.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
-            whenever(it.transportInfo).thenReturn(wifiInfo)
+            whenever(it.transportInfo).thenReturn(transportInfo)
             whenever(it.hasCapability(NET_CAPABILITY_VALIDATED)).thenReturn(isValidated)
         }
     }