From ecd0886e3619b3b7ccaf7a5b971403aa076dd0fb Mon Sep 17 00:00:00 2001
From: Chiachang Wang <chiachangwang@google.com>
Date: Fri, 21 May 2021 08:51:58 +0000
Subject: [PATCH] Unify the verification for unregister a NetworkAgent

It's a follow up of aosp/1652210.

The tests may verify partial NetworkAgent callback behavior.
The existing expectCallback will always expect to see the first
callback in the callback queue, so if tests did not verify all
the happened callback. It cannot verify the callback after
unregister. Add eventuallyExpect to wait for a target callback
and skip the disinterest ones, and verify the unregister
behavior for NetworkAgent callback in one method.

Bug: 178725261
Test: atest android.net.cts.NetworkAgentTest
Change-Id: I66d8e5a0fa1e2245711e8ac90d9daca24802e399
Merged-In: I66d8e5a0fa1e2245711e8ac90d9daca24802e399
(cherry picked from commit 12a6422f47d98723eb74e25b170185908ca8947d)
---
 .../src/android/net/cts/NetworkAgentTest.kt   | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 1c9aba1e11..8e2b3101ed 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 @@ class NetworkAgentTest {
             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 @@ class NetworkAgentTest {
         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 @@ class NetworkAgentTest {
         agent.expectNoInternetValidationStatus()
         mCM.requestBandwidthUpdate(agent.network)
         agent.expectCallback<OnBandwidthUpdateRequested>()
-        agent.unregister()
+        unregister(agent)
     }
 
     @Test
@@ -648,10 +652,16 @@ class NetworkAgentTest {
             }
         }
 
-        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() {