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
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 51a4f32..16c7477 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 @@
 
     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 @@
             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 @@
             // 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();