From f9093389f5f84c30ed490cf9416fec57ebfe24d8 Mon Sep 17 00:00:00 2001
From: Chiachang Wang <chiachangwang@google.com>
Date: Tue, 6 Apr 2021 20:21:34 +0800
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
---
 .../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 f53a2a80b3..865a07a52a 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -300,6 +300,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")
@@ -376,13 +381,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
@@ -393,7 +397,7 @@ class NetworkAgentTest {
         agent.expectNoInternetValidationStatus()
         mCM.requestBandwidthUpdate(agent.network)
         agent.expectCallback<OnBandwidthUpdateRequested>()
-        agent.unregister()
+        unregister(agent)
     }
 
     @Test
@@ -635,10 +639,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() {