Fix testIpClientServiceAvailableAfterPolling am: a47091526a

Change-Id: Icca83bd6064c263ab352fc7d0bcacb18d9d24574

1 files changed, 10 insertions, 1 deletions

diff --git a/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt b/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
index 2d6a095..31b4ffd 100644
--- a/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
+++ b/tests/unit/src/android/net/networkstack/ModuleNetworkStackClientTest.kt
@@ -29,6 +29,8 @@ import android.os.Build
 import android.os.IBinder
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.testutils.DevSdkIgnoreRule
 import org.junit.After
 import org.junit.Before
@@ -37,7 +39,6 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.any
 import org.mockito.Mock
-import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.never
 import org.mockito.Mockito.timeout
 import org.mockito.Mockito.verify
@@ -91,11 +92,19 @@ class ModuleNetworkStackClientTest {
 
     @Test
     fun testIpClientServiceAvailableAfterPolling() {
+        // Force NetworkStack.getService() to return null: this cannot be done with
+        // setServiceForTest, as passing null just restores default behavior.
+        val session = mockitoSession().spyStatic(NetworkStack::class.java).startMocking()
+        doReturn(null).`when` { NetworkStack.getService() }
         ModuleNetworkStackClient.getInstance(mContext).makeIpClient(TEST_IFNAME, mIpClientCb)
 
         Thread.sleep(TEST_TIMEOUT_MS)
         verify(mConnector, never()).makeIpClient(any(), any())
         NetworkStack.setServiceForTest(mConnectorBinder)
+
+        // Restore behavior of NetworkStack to return what was set in setServiceForTest
+        session.finishMocking()
+
         // Use a longer timeout as polling can cause larger delays
         verify(mConnector, timeout(TEST_TIMEOUT_MS * 4)).makeIpClient(TEST_IFNAME, mIpClientCb)
     }