Make TestableNetworkCallback understand blocked reasons. am: 5d4e7a5d36 am: b14d6f3014

Original change: https://android-review.googlesource.com/c/platform/frameworks/libs/net/+/1652229

Change-Id: I90fc1bd0b46d63e8dd0c4235ffa6241c7ec95aaa
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt b/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt
index 043654f..1b3d0f6 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt
+++ b/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt
@@ -25,6 +25,7 @@
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.testutils.RecorderCallback.CallbackEntry.Available
 import com.android.testutils.RecorderCallback.CallbackEntry.BlockedStatus
+import com.android.testutils.RecorderCallback.CallbackEntry.BlockedStatusInt
 import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
 import com.android.testutils.RecorderCallback.CallbackEntry.Losing
@@ -82,7 +83,10 @@
             override val network: Network,
             val blocked: Boolean
         ) : CallbackEntry()
-
+        data class BlockedStatusInt(
+            override val network: Network,
+            val blocked: Int
+        ) : CallbackEntry()
         // Convenience constants for expecting a type
         companion object {
             @JvmField
@@ -103,6 +107,8 @@
             val UNAVAILABLE = Unavailable::class
             @JvmField
             val BLOCKED_STATUS = BlockedStatus::class
+            @JvmField
+            val BLOCKED_STATUS_INT = BlockedStatusInt::class
         }
     }
 
@@ -131,6 +137,9 @@
         history.add(BlockedStatus(network, blocked))
     }
 
+    // Cannot do:
+    // fun onBlockedStatusChanged(network: Network, blocked: Int) {
+    // because on S, that needs to be "override fun", and on R, that cannot be "override fun".
     override fun onNetworkSuspended(network: Network) {
         Log.d(TAG, "onNetworkSuspended $network $network")
         history.add(Suspended(network))
@@ -272,13 +281,33 @@
         blocked: Boolean = false,
         tmt: Long = defaultTimeoutMs
     ) {
+        expectAvailableCallbacksCommon(net, suspended, validated, tmt)
+        expectBlockedStatusCallback(blocked, net, tmt)
+    }
+
+    fun expectAvailableCallbacks(
+        net: Network,
+        suspended: Boolean,
+        validated: Boolean,
+        blockedStatus: Int,
+        tmt: Long
+    ) {
+        expectAvailableCallbacksCommon(net, suspended, validated, tmt)
+        expectBlockedStatusCallback(blockedStatus, net)
+    }
+
+    private fun expectAvailableCallbacksCommon(
+        net: Network,
+        suspended: Boolean,
+        validated: Boolean,
+        tmt: Long
+    ) {
         expectCallback<Available>(net, tmt)
         if (suspended) {
             expectCallback<Suspended>(net, tmt)
         }
         expectCapabilitiesThat(net, tmt) { validated == it.hasCapability(NET_CAPABILITY_VALIDATED) }
         expectCallback<LinkPropertiesChanged>(net, tmt)
-        expectBlockedStatusCallback(blocked, net)
     }
 
     // Backward compatibility for existing Java code. Use named arguments instead and remove all
@@ -295,6 +324,12 @@
         }
     }
 
+    fun expectBlockedStatusCallback(blocked: Int, net: Network, tmt: Long = defaultTimeoutMs) {
+        expectCallback<BlockedStatusInt>(net, tmt).also {
+            assertEquals(it.blocked, blocked, "Unexpected blocked status ${it.blocked}")
+        }
+    }
+
     // Expects the available callbacks (where the onCapabilitiesChanged must contain the
     // VALIDATED capability), plus another onCapabilitiesChanged which is identical to the
     // one we just sent.
@@ -314,6 +349,16 @@
         expectCapabilitiesThat(net, tmt) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
     }
 
+    fun expectAvailableThenValidatedCallbacks(
+        net: Network,
+        blockedStatus: Int,
+        tmt: Long = defaultTimeoutMs
+    ) {
+        expectAvailableCallbacks(net, validated = false, suspended = false,
+                blockedStatus = blockedStatus, tmt = tmt)
+        expectCapabilitiesThat(net, tmt) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
+    }
+
     // Temporary Java compat measure : have MockNetworkAgent implement this so that all existing
     // calls with networkAgent can be routed through here without moving MockNetworkAgent.
     // TODO: clean this up, remove this method.