Merge "[SB Refactor] Add table logging to the wifi view model." into tm-qpr-dev am: 49480248da am: 3e5a078ee3

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/20588810

Change-Id: I637edb5d235ac1b38a3b12f4cb7d13f7640e8df0
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

11 files changed, 142 insertions, 66 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
index b816364ed4cf..52237605caf9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
@@ -24,6 +24,7 @@ import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.statusbar.phone.StatusBarIconController
 import com.android.systemui.statusbar.phone.StatusBarLocation
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.WifiViewModel
 import javax.inject.Inject
@@ -73,7 +74,9 @@ constructor(
                         // Note that this flow may still run if
                         // [statusBarPipelineFlags.runNewWifiIconBackend] is true because we may
                         // want to get the logging data without rendering.
-                        if (wifiIcon != null && statusBarPipelineFlags.useNewWifiIcon()) {
+                        if (
+                            wifiIcon is WifiIcon.Visible && statusBarPipelineFlags.useNewWifiIcon()
+                        ) {
                             iconController.setNewWifiIcon()
                         }
                     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
index 345f8cb75660..f5b5950d33a8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
@@ -30,6 +30,7 @@ import com.android.systemui.statusbar.StatusBarIconView
 import com.android.systemui.statusbar.StatusBarIconView.STATE_DOT
 import com.android.systemui.statusbar.StatusBarIconView.STATE_HIDDEN
 import com.android.systemui.statusbar.StatusBarIconView.STATE_ICON
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel
 import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -92,8 +93,10 @@ object WifiViewBinder {
 
                 launch {
                     viewModel.wifiIcon.collect { wifiIcon ->
-                        view.isVisible = wifiIcon != null
-                        wifiIcon?.let { IconViewBinder.bind(wifiIcon, iconView) }
+                        view.isVisible = wifiIcon is WifiIcon.Visible
+                        if (wifiIcon is WifiIcon.Visible) {
+                            IconViewBinder.bind(wifiIcon.icon, iconView)
+                        }
                     }
                 }
 
@@ -135,7 +138,7 @@ object WifiViewBinder {
 
         return object : Binding {
             override fun getShouldIconBeVisible(): Boolean {
-                return viewModel.wifiIcon.value != null
+                return viewModel.wifiIcon.value is WifiIcon.Visible
             }
 
             override fun onVisibilityStateChanged(@StatusBarIconView.VisibleState state: Int) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt
new file mode 100644
index 000000000000..e491d2bbf0d6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.pipeline.wifi.ui.model
+
+import android.annotation.DrawableRes
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.log.table.Diffable
+import com.android.systemui.log.table.TableRowLogger
+
+/** Represents the various states of the wifi icon. */
+sealed interface WifiIcon : Diffable<WifiIcon> {
+    /** Represents a wifi icon that should be hidden (not visible). */
+    object Hidden : WifiIcon {
+        override fun toString() = "hidden"
+    }
+
+    /**
+     * Represents a visible wifi icon that uses [res] as its image and [contentDescription] as its
+     * description.
+     */
+    class Visible(
+        @DrawableRes res: Int,
+        val contentDescription: ContentDescription.Loaded,
+    ) : WifiIcon {
+        val icon = Icon.Resource(res, contentDescription)
+
+        override fun toString() = contentDescription.description.toString()
+    }
+
+    override fun logDiffs(prevVal: WifiIcon, row: TableRowLogger) {
+        if (prevVal.toString() != toString()) {
+            row.logChange(COL_ICON, toString())
+        }
+    }
+
+    override fun logFull(row: TableRowLogger) {
+        row.logChange(COL_ICON, toString())
+    }
+}
+
+private const val COL_ICON = "wifiIcon"
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
index 95ab251422b2..a29c9b94e6b8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
@@ -28,7 +28,7 @@ import kotlinx.coroutines.flow.StateFlow
  */
 class HomeWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
index 86535d63f84f..1e190fb898ff 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
@@ -17,15 +17,15 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /** A view model for the wifi icon shown on keyguard (lockscreen). */
 class KeyguardWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
index 7cbdf5dbdf2d..e35a8fef4528 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.flowOf
@@ -33,8 +33,8 @@ abstract class LocationBasedWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
     debugTint: Int,
 
-    /** The wifi icon that should be displayed. Null if we shouldn't display any icon. */
-    val wifiIcon: StateFlow<Icon.Resource?>,
+    /** The wifi icon that should be displayed. */
+    val wifiIcon: StateFlow<WifiIcon>,
 
     /** True if the activity in view should be visible. */
     val isActivityInViewVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
index fd54c5f5062e..18e62b284cb9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
@@ -17,15 +17,15 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /** A view model for the wifi icon shown in quick settings (when the shade is pulled down). */
 class QsWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
index 0782bbb774eb..ec7ba653cac7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
@@ -17,16 +17,18 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.content.Context
-import androidx.annotation.DrawableRes
 import androidx.annotation.StringRes
 import androidx.annotation.VisibleForTesting
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_CONNECTION_STRENGTH
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_NO_CONNECTION
 import com.android.systemui.R
 import com.android.systemui.common.shared.model.ContentDescription
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.pipeline.dagger.WifiTableLog
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
+import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logDiffsForTable
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_FULL_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_INTERNET_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_NETWORK
@@ -71,50 +73,39 @@ constructor(
     connectivityConstants: ConnectivityConstants,
     private val context: Context,
     logger: ConnectivityPipelineLogger,
+    @WifiTableLog wifiTableLogBuffer: TableLogBuffer,
     interactor: WifiInteractor,
     @Application private val scope: CoroutineScope,
     statusBarPipelineFlags: StatusBarPipelineFlags,
     wifiConstants: WifiConstants,
 ) {
-    /**
-     * Returns the drawable resource ID to use for the wifi icon based on the given network.
-     * Null if we can't compute the icon.
-     */
-    @DrawableRes
-    private fun WifiNetworkModel.iconResId(): Int? {
+    /** Returns the icon to use based on the given network. */
+    private fun WifiNetworkModel.icon(): WifiIcon {
         return when (this) {
-            is WifiNetworkModel.CarrierMerged -> null
-            is WifiNetworkModel.Inactive -> WIFI_NO_NETWORK
-            is WifiNetworkModel.Active ->
-                when {
-                    this.level == null -> null
-                    this.isValidated -> WIFI_FULL_ICONS[this.level]
-                    else -> WIFI_NO_INTERNET_ICONS[this.level]
-                }
-        }
-    }
-
-    /**
-     * Returns the content description for the wifi icon based on the given network.
-     * Null if we can't compute the content description.
-     */
-    private fun WifiNetworkModel.contentDescription(): ContentDescription? {
-        return when (this) {
-            is WifiNetworkModel.CarrierMerged -> null
-            is WifiNetworkModel.Inactive ->
+            is WifiNetworkModel.CarrierMerged -> WifiIcon.Hidden
+            is WifiNetworkModel.Inactive -> WifiIcon.Visible(
+                res = WIFI_NO_NETWORK,
                 ContentDescription.Loaded(
                     "${context.getString(WIFI_NO_CONNECTION)},${context.getString(NO_INTERNET)}"
                 )
+            )
             is WifiNetworkModel.Active ->
                 when (this.level) {
-                    null -> null
+                    null -> WifiIcon.Hidden
                     else -> {
                         val levelDesc = context.getString(WIFI_CONNECTION_STRENGTH[this.level])
                         when {
-                            this.isValidated -> ContentDescription.Loaded(levelDesc)
+                            this.isValidated ->
+                                WifiIcon.Visible(
+                                    WIFI_FULL_ICONS[this.level],
+                                    ContentDescription.Loaded(levelDesc)
+                                )
                             else ->
-                                ContentDescription.Loaded(
-                                    "$levelDesc,${context.getString(NO_INTERNET)}"
+                                WifiIcon.Visible(
+                                    WIFI_NO_INTERNET_ICONS[this.level],
+                                    ContentDescription.Loaded(
+                                        "$levelDesc,${context.getString(NO_INTERNET)}"
+                                    )
                                 )
                         }
                     }
@@ -122,8 +113,8 @@ constructor(
         }
     }
 
-    /** The wifi icon that should be displayed. Null if we shouldn't display any icon. */
-    private val wifiIcon: StateFlow<Icon.Resource?> =
+    /** The wifi icon that should be displayed. */
+    private val wifiIcon: StateFlow<WifiIcon> =
         combine(
             interactor.isEnabled,
             interactor.isDefault,
@@ -131,22 +122,29 @@ constructor(
             interactor.wifiNetwork,
         ) { isEnabled, isDefault, isForceHidden, wifiNetwork ->
             if (!isEnabled || isForceHidden || wifiNetwork is WifiNetworkModel.CarrierMerged) {
-                return@combine null
+                return@combine WifiIcon.Hidden
             }
 
-            val iconResId = wifiNetwork.iconResId() ?: return@combine null
-            val icon = Icon.Resource(iconResId, wifiNetwork.contentDescription())
+            val icon = wifiNetwork.icon()
 
             return@combine when {
                 isDefault -> icon
                 wifiConstants.alwaysShowIconIfEnabled -> icon
                 !connectivityConstants.hasDataCapabilities -> icon
                 wifiNetwork is WifiNetworkModel.Active && wifiNetwork.isValidated -> icon
-                else -> null
+                else -> WifiIcon.Hidden
             }
         }
-                .logOutputChange(logger, "icon") { icon -> icon?.contentDescription.toString() }
-                .stateIn(scope, started = SharingStarted.WhileSubscribed(), initialValue = null)
+            .logDiffsForTable(
+                wifiTableLogBuffer,
+                columnPrefix = "",
+                initialValue = WifiIcon.Hidden,
+            )
+            .stateIn(
+                scope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = WifiIcon.Hidden
+            )
 
     /** The wifi activity status. Null if we shouldn't display the activity status. */
     private val activity: Flow<WifiActivityModel?> =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
index 5c16e1295b65..3d9fd961222f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
@@ -25,6 +25,7 @@ import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.lifecycle.InstantTaskExecutorRule
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.StatusBarIconView.STATE_DOT
 import com.android.systemui.statusbar.StatusBarIconView.STATE_HIDDEN
 import com.android.systemui.statusbar.StatusBarIconView.STATE_ICON
@@ -64,6 +65,7 @@ class ModernStatusBarWifiViewTest : SysuiTestCase() {
     private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock
     private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock
     private lateinit var connectivityConstants: ConnectivityConstants
     @Mock
@@ -103,6 +105,7 @@ class ModernStatusBarWifiViewTest : SysuiTestCase() {
             connectivityConstants,
             context,
             logger,
+            tableLogBuffer,
             interactor,
             scope,
             statusBarPipelineFlags,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
index 3001b8162185..12b93819fc5e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
@@ -23,6 +23,7 @@ import com.android.settingslib.AccessibilityContentDescriptions.WIFI_CONNECTION_
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_NO_CONNECTION
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.ContentDescription.Companion.loadContentDescription
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_FULL_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_INTERNET_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_NETWORK
@@ -40,6 +41,7 @@ import com.android.systemui.statusbar.pipeline.wifi.data.repository.FakeWifiRepo
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractor
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractorImpl
 import com.android.systemui.statusbar.pipeline.wifi.shared.WifiConstants
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.WifiViewModel.Companion.NO_INTERNET
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
@@ -67,6 +69,7 @@ internal class WifiViewModelIconParameterizedTest(private val testCase: TestCase
 
     @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -123,6 +126,7 @@ internal class WifiViewModelIconParameterizedTest(private val testCase: TestCase
                     connectivityConstants,
                     context,
                     logger,
+                    tableLogBuffer,
                     interactor,
                     scope,
                     statusBarPipelineFlags,
@@ -137,15 +141,21 @@ internal class WifiViewModelIconParameterizedTest(private val testCase: TestCase
             yield()
 
             // THEN we get the expected icon
-            assertThat(iconFlow.value?.res).isEqualTo(testCase.expected?.iconResource)
-            val expectedContentDescription =
-                if (testCase.expected == null) {
-                    null
-                } else {
-                    testCase.expected.contentDescription.invoke(context)
+            val actualIcon = iconFlow.value
+            when (testCase.expected) {
+                null -> {
+                    assertThat(actualIcon).isInstanceOf(WifiIcon.Hidden::class.java)
+                }
+                else -> {
+                    assertThat(actualIcon).isInstanceOf(WifiIcon.Visible::class.java)
+                    val actualIconVisible = actualIcon as WifiIcon.Visible
+                    assertThat(actualIconVisible.icon.res).isEqualTo(testCase.expected.iconResource)
+                    val expectedContentDescription =
+                        testCase.expected.contentDescription.invoke(context)
+                    assertThat(actualIconVisible.contentDescription.loadContentDescription(context))
+                        .isEqualTo(expectedContentDescription)
                 }
-            assertThat(iconFlow.value?.contentDescription?.loadContentDescription(context))
-                .isEqualTo(expectedContentDescription)
+            }
 
             job.cancel()
         }
@@ -174,7 +184,7 @@ internal class WifiViewModelIconParameterizedTest(private val testCase: TestCase
         val isDefault: Boolean = false,
         val network: WifiNetworkModel,
 
-        /** The expected output. Null if we expect the output to be null. */
+        /** The expected output. Null if we expect the output to be hidden. */
         val expected: Expected?
     ) {
         override fun toString(): String {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
index 6a6b2a801ab0..7502020286f0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
@@ -18,7 +18,7 @@ package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
@@ -34,6 +34,7 @@ import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiIntera
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractorImpl
 import com.android.systemui.statusbar.pipeline.wifi.shared.WifiConstants
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiActivityModel
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -59,6 +60,7 @@ class WifiViewModelTest : SysuiTestCase() {
 
     @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -103,21 +105,21 @@ class WifiViewModelTest : SysuiTestCase() {
 
     @Test
     fun wifiIcon_allLocationViewModelsReceiveSameData() = runBlocking(IMMEDIATE) {
-        var latestHome: Icon? = null
+        var latestHome: WifiIcon? = null
         val jobHome = underTest
             .home
             .wifiIcon
             .onEach { latestHome = it }
             .launchIn(this)
 
-        var latestKeyguard: Icon? = null
+        var latestKeyguard: WifiIcon? = null
         val jobKeyguard = underTest
             .keyguard
             .wifiIcon
             .onEach { latestKeyguard = it }
             .launchIn(this)
 
-        var latestQs: Icon? = null
+        var latestQs: WifiIcon? = null
         val jobQs = underTest
             .qs
             .wifiIcon
@@ -133,7 +135,7 @@ class WifiViewModelTest : SysuiTestCase() {
         )
         yield()
 
-        assertThat(latestHome).isInstanceOf(Icon.Resource::class.java)
+        assertThat(latestHome).isInstanceOf(WifiIcon.Visible::class.java)
         assertThat(latestHome).isEqualTo(latestKeyguard)
         assertThat(latestKeyguard).isEqualTo(latestQs)
 
@@ -541,6 +543,7 @@ class WifiViewModelTest : SysuiTestCase() {
             connectivityConstants,
             context,
             logger,
+            tableLogBuffer,
             interactor,
             scope,
             statusBarPipelineFlags,