From 4da74f040ac8891b2bcd2b8b2332ac4f9037ee41 Mon Sep 17 00:00:00 2001
From: Roshan Pius <rpius@google.com>
Date: Thu, 19 Mar 2020 06:33:10 -0700
Subject: [PATCH] WifiManagerTest: Make traffic state change test more robust

Traffic state changed callbacks can be triggered multiple times with
different data indicators. Ensure that we wait for the expected state
before unblocking.

Bug: 151900245
Test: atest android.net.wifi.cts.WifiManagerTest
Change-Id: Ie451070bf2ab430f57ff83b0cefca7050aa8df59
---
 .../net/src/android/net/wifi/cts/WifiManagerTest.java | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
index 51a4f32a64..16c7477aff 100644
--- a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
@@ -1840,11 +1840,13 @@ public class WifiManagerTest extends AndroidTestCase {
 
     private static class TestTrafficStateCallback implements WifiManager.TrafficStateCallback {
         private final Object mLock;
+        private final int mWaitForState;
         public boolean onStateChangedCalled = false;
         public int state = -1;
 
-        TestTrafficStateCallback(Object lock) {
+        TestTrafficStateCallback(Object lock, int waitForState) {
             mLock = lock;
+            mWaitForState = waitForState;
         }
 
         @Override
@@ -1852,7 +1854,9 @@ public class WifiManagerTest extends AndroidTestCase {
             synchronized (mLock) {
                 onStateChangedCalled = true;
                 this.state = state;
-                mLock.notify();
+                if (mWaitForState == state) { // only notify if we got the expected state.
+                    mLock.notify();
+                }
             }
         }
     }
@@ -1886,7 +1890,8 @@ public class WifiManagerTest extends AndroidTestCase {
             // skip the test if WiFi is not supported
             return;
         }
-        TestTrafficStateCallback trafficStateCallback = new TestTrafficStateCallback(mLock);
+        TestTrafficStateCallback trafficStateCallback =
+                new TestTrafficStateCallback(mLock, DATA_ACTIVITY_INOUT);
         UiAutomation uiAutomation = InstrumentationRegistry.getInstrumentation().getUiAutomation();
         try {
             uiAutomation.adoptShellPermissionIdentity();