diff --git a/staticlibs/tests/unit/src/com/android/testutils/TestableNetworkCallbackTest.kt b/staticlibs/tests/unit/src/com/android/testutils/TestableNetworkCallbackTest.kt
index 557f1fd..690ef2e 100644
--- a/staticlibs/tests/unit/src/com/android/testutils/TestableNetworkCallbackTest.kt
+++ b/staticlibs/tests/unit/src/com/android/testutils/TestableNetworkCallbackTest.kt
@@ -270,18 +270,30 @@
     }
 
     @Test
-    fun testPollForNextCallback() {
-        assertFails { mCallback.pollForNextCallback(SHORT_TIMEOUT_MS) }
+    fun testPoll() {
+        assertNull(mCallback.poll(SHORT_TIMEOUT_MS))
         TNCInterpreter.interpretTestSpec(initial = mCallback, lineShift = 1,
                 threadTransform = { cb -> cb.createLinkedCopy() }, spec = """
             sleep; onAvailable(133)    | poll(2) = Available(133) time 1..4
-                                       | poll(1) fails
+                                       | poll(1) = null
             onCapabilitiesChanged(108) | poll(1) = CapabilitiesChanged(108) time 0..3
             onBlockedStatus(199)       | poll(1) = BlockedStatus(199) time 0..3
         """)
     }
 
     @Test
+    fun testPollOrThrow() {
+        assertFails { mCallback.pollOrThrow(SHORT_TIMEOUT_MS) }
+        TNCInterpreter.interpretTestSpec(initial = mCallback, lineShift = 1,
+                threadTransform = { cb -> cb.createLinkedCopy() }, spec = """
+            sleep; onAvailable(133)    | pollOrThrow(2) = Available(133) time 1..4
+                                       | pollOrThrow(1) fails
+            onCapabilitiesChanged(108) | pollOrThrow(1) = CapabilitiesChanged(108) time 0..3
+            onBlockedStatus(199)       | pollOrThrow(1) = BlockedStatus(199) time 0..3
+        """)
+    }
+
+    @Test
     fun testEventuallyExpect() {
         // TODO: Current test does not verify the inline one. Also verify the behavior after
         // aligning two eventuallyExpect()
@@ -358,9 +370,8 @@
             else -> fail("Unknown callback type")
         }
     },
-    Regex("""poll\((\d+)\)""") to { i, cb, t ->
-        cb.pollForNextCallback(t.timeArg(1))
-    },
+    Regex("""poll\((\d+)\)""") to { i, cb, t -> cb.poll(t.timeArg(1)) },
+    Regex("""pollOrThrow\((\d+)\)""") to { i, cb, t -> cb.pollOrThrow(t.timeArg(1)) },
     // Interpret "eventually(Available(xx), timeout)" as calling eventuallyExpect that expects
     // CallbackEntry.AVAILABLE with netId of xx within timeout*INTERPRET_TIME_UNIT timeout, and
     // likewise for all callback types.
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
index 4407714..0c16937 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
@@ -43,8 +43,6 @@
 object ANY_NETWORK : Network(-2)
 fun anyNetwork() = ANY_NETWORK
 
-private val Int.capabilityName get() = NetworkCapabilities.capabilityNameOf(this)
-
 open class RecorderCallback private constructor(
     private val backingRecord: ArrayTrackRecord<CallbackEntry>
 ) : NetworkCallback() {
@@ -193,9 +191,26 @@
             else -> null
         }
 
-    fun pollForNextCallback(timeoutMs: Long = defaultTimeoutMs): CallbackEntry {
-        return history.poll(timeoutMs) ?: fail("Did not receive callback after ${timeoutMs}ms")
-    }
+    /**
+     * Get the next callback or null if timeout.
+     *
+     * With no argument, this method waits out the default timeout. To wait forever, pass
+     * Long.MAX_VALUE.
+     */
+    @JvmOverloads
+    fun poll(timeoutMs: Long = defaultTimeoutMs): CallbackEntry? = history.poll(timeoutMs)
+
+    /**
+     * Get the next callback or throw if timeout.
+     *
+     * With no argument, this method waits out the default timeout. To wait forever, pass
+     * Long.MAX_VALUE.
+     */
+    @JvmOverloads
+    fun pollOrThrow(
+        timeoutMs: Long = defaultTimeoutMs,
+        errorMsg: String = "Did not receive callback after $timeoutMs"
+    ): CallbackEntry = poll(timeoutMs) ?: fail(errorMsg)
 
     // Make open for use in ConnectivityServiceTest which is the only one knowing its handlers.
     // TODO : remove the necessity to overload this, remove the open qualifier, and give a
@@ -221,7 +236,7 @@
     inline fun <reified T : CallbackEntry> expectCallback(
         network: Network = ANY_NETWORK,
         timeoutMs: Long = defaultTimeoutMs
-    ): T = pollForNextCallback(timeoutMs).let {
+    ): T = pollOrThrow(timeoutMs).let {
         if (it !is T || (ANY_NETWORK !== network && it.network != network)) {
             fail("Unexpected callback : $it, expected ${T::class} with Network[$network]")
         } else {
@@ -258,7 +273,7 @@
     fun expectCallbackThat(
         timeoutMs: Long = defaultTimeoutMs,
         valid: (CallbackEntry) -> Boolean
-    ) = pollForNextCallback(timeoutMs).also { assertTrue(valid(it), "Unexpected callback : $it") }
+    ) = pollOrThrow(timeoutMs).also { assertTrue(valid(it), "Unexpected callback : $it") }
 
     fun expectCapabilitiesThat(
         net: Network,
@@ -392,7 +407,7 @@
         type: KClass<T>,
         n: Network?,
         timeoutMs: Long = defaultTimeoutMs
-    ) = pollForNextCallback(timeoutMs).also {
+    ) = pollOrThrow(timeoutMs).also {
         val network = n ?: NULL_NETWORK
         // TODO : remove this .java access if the tests ever use kotlin-reflect. At the time of
         // this writing this would be the only use of this library in the tests.