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 65d9dea41e..ae5ef75af6 100644
--- a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
@@ -1956,11 +1956,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
@@ -1968,7 +1970,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();
+                }
             }
         }
     }
@@ -2002,7 +2006,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();