diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 1c9aba1..8e2b310 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -306,6 +306,11 @@
             return foundCallback
         }
 
+        inline fun <reified T : CallbackEntry> eventuallyExpect() =
+                history.poll(DEFAULT_TIMEOUT_MS) { it is T }.also {
+                    assertNotNull(it, "Callback ${T::class} not received")
+        } as T
+
         fun assertNoCallback() {
             assertTrue(waitForIdle(DEFAULT_TIMEOUT_MS),
                     "Handler didn't became idle after ${DEFAULT_TIMEOUT_MS}ms")
@@ -383,13 +388,12 @@
         callback.expectAvailableThenValidatedCallbacks(agent.network)
         agent.expectEmptySignalStrengths()
         agent.expectNoInternetValidationStatus()
-        agent.unregister()
+
+        unregister(agent)
         callback.expectCallback<Lost>(agent.network)
-        agent.expectCallback<OnNetworkUnwanted>()
         assertFailsWith<IllegalStateException>("Must not be able to register an agent twice") {
             agent.register()
         }
-        agent.expectCallback<OnNetworkDestroyed>()
     }
 
     @Test
@@ -400,7 +404,7 @@
         agent.expectNoInternetValidationStatus()
         mCM.requestBandwidthUpdate(agent.network)
         agent.expectCallback<OnBandwidthUpdateRequested>()
-        agent.unregister()
+        unregister(agent)
     }
 
     @Test
@@ -648,10 +652,16 @@
             }
         }
 
-        agent.unregister()
+        unregister(agent)
         callback.expectCallback<Lost>(agent.network)
     }
 
+    private fun unregister(agent: TestableNetworkAgent) {
+        agent.unregister()
+        agent.eventuallyExpect<OnNetworkUnwanted>()
+        agent.eventuallyExpect<OnNetworkDestroyed>()
+    }
+
     @Test
     @IgnoreUpTo(Build.VERSION_CODES.R)
     fun testAgentStartsInConnecting() {