Merge "WifiServiceImpl: Add callback for traffic poller"

4 files changed, 260 insertions, 87 deletions

diff --git a/service/java/com/android/server/wifi/WifiServiceImpl.java b/service/java/com/android/server/wifi/WifiServiceImpl.java
index 30fd29514..a805f4a49 100644
--- a/service/java/com/android/server/wifi/WifiServiceImpl.java
+++ b/service/java/com/android/server/wifi/WifiServiceImpl.java
@@ -58,6 +58,7 @@ import android.net.NetworkUtils;
 import android.net.Uri;
 import android.net.ip.IpClient;
 import android.net.wifi.ISoftApCallback;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.IWifiManager;
 import android.net.wifi.ScanResult;
 import android.net.wifi.WifiActivityEnergyInfo;
@@ -250,26 +251,6 @@ public class WifiServiceImpl extends IWifiManager.Stub {
         public void handleMessage(Message msg) {
             super.handleMessage(msg);
             switch (msg.what) {
-                case AsyncChannel.CMD_CHANNEL_HALF_CONNECTED: {
-                    if (msg.arg1 == AsyncChannel.STATUS_SUCCESSFUL) {
-                        Slog.d(TAG, "New client listening to asynchronous messages");
-                        // We track the clients by the Messenger
-                        // since it is expected to be always available
-                        mTrafficPoller.addClient(msg.replyTo);
-                    } else {
-                        Slog.e(TAG, "Client connection failure, error=" + msg.arg1);
-                    }
-                    break;
-                }
-                case AsyncChannel.CMD_CHANNEL_DISCONNECTED: {
-                    if (msg.arg1 == AsyncChannel.STATUS_SEND_UNSUCCESSFUL) {
-                        Slog.w(TAG, "Send failed, client connection lost");
-                    } else {
-                        Slog.w(TAG, "Client connection lost with reason: " + msg.arg1);
-                    }
-                    mTrafficPoller.removeClient(msg.replyTo);
-                    break;
-                }
                 case AsyncChannel.CMD_CHANNEL_FULL_CONNECTION: {
                     AsyncChannel ac = mFrameworkFacade.makeWifiAsyncChannel(TAG);
                     ac.connect(mContext, this, msg.replyTo);
@@ -2862,4 +2843,75 @@ public class WifiServiceImpl extends IWifiManager.Stub {
                     .c(provider.toString()).flush();
         }
     }
+
+    /**
+     * see {@link android.net.wifi.WifiManager#registerTrafficStateCallback(
+     * TrafficStateCallback, Handler)}
+     *
+     * @param binder IBinder instance to allow cleanup if the app dies
+     * @param callback Traffic State callback to register
+     * @param callbackIdentifier Unique ID of the registering callback. This ID will be used to
+     *        unregister the callback. See {@link unregisterTrafficStateCallback(int)}
+     *
+     * @throws SecurityException if the caller does not have permission to register a callback
+     * @throws RemoteException if remote exception happens
+     * @throws IllegalArgumentException if the arguments are null or invalid
+     */
+    @Override
+    public void registerTrafficStateCallback(IBinder binder, ITrafficStateCallback callback,
+                                             int callbackIdentifier) {
+        // verify arguments
+        if (binder == null) {
+            throw new IllegalArgumentException("Binder must not be null");
+        }
+        if (callback == null) {
+            throw new IllegalArgumentException("Callback must not be null");
+        }
+        enforceNetworkSettingsPermission();
+        if (mVerboseLoggingEnabled) {
+            mLog.info("registerTrafficStateCallback uid=%").c(Binder.getCallingUid()).flush();
+        }
+
+        // register for binder death
+        IBinder.DeathRecipient dr = new IBinder.DeathRecipient() {
+            @Override
+            public void binderDied() {
+                binder.unlinkToDeath(this, 0);
+                mClientHandler.post(() -> {
+                    mTrafficPoller.removeCallback(callbackIdentifier);
+                });
+            }
+        };
+        try {
+            binder.linkToDeath(dr, 0);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error on linkToDeath - " + e);
+            return;
+        }
+        // Post operation to handler thread
+        mClientHandler.post(() -> {
+            mTrafficPoller.addCallback(callback, callbackIdentifier);
+        });
+    }
+
+    /**
+     * see {@link android.net.wifi.WifiManager#unregisterTrafficStateCallback(
+     * WifiManager.TrafficStateCallback)}
+     *
+     * @param callbackIdentifier Unique ID of the callback to be unregistered.
+     *
+     * @throws SecurityException if the caller does not have permission to register a callback
+     */
+    @Override
+    public void unregisterTrafficStateCallback(int callbackIdentifier) {
+        enforceNetworkSettingsPermission();
+        if (mVerboseLoggingEnabled) {
+            mLog.info("unregisterTrafficStateCallback uid=%").c(Binder.getCallingUid()).flush();
+        }
+
+        // Post operation to handler thread
+        mClientHandler.post(() -> {
+            mTrafficPoller.removeCallback(callbackIdentifier);
+        });
+    }
 }
diff --git a/service/java/com/android/server/wifi/WifiTrafficPoller.java b/service/java/com/android/server/wifi/WifiTrafficPoller.java
index bcb625ed6..d7f02db89 100644
--- a/service/java/com/android/server/wifi/WifiTrafficPoller.java
+++ b/service/java/com/android/server/wifi/WifiTrafficPoller.java
@@ -24,19 +24,18 @@ import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.net.NetworkInfo;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
-import android.os.Messenger;
 import android.os.RemoteException;
 import android.text.TextUtils;
 import android.util.Log;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -51,11 +50,12 @@ public class WifiTrafficPoller {
      * statistics
      */
     private static final int POLL_TRAFFIC_STATS_INTERVAL_MSECS = 1000;
+    /* Limit on number of registered soft AP callbacks to track and prevent potential memory leak */
+    private static final int NUM_CALLBACKS_WARN_LIMIT = 10;
+    private static final int NUM_CALLBACKS_WTF_LIMIT = 20;
 
     private static final int ENABLE_TRAFFIC_STATS_POLL  = 1;
     private static final int TRAFFIC_STATS_POLL         = 2;
-    private static final int ADD_CLIENT                 = 3;
-    private static final int REMOVE_CLIENT              = 4;
 
     private boolean mEnableTrafficStatsPoll = false;
     private int mTrafficStatsPollToken = 0;
@@ -64,7 +64,7 @@ public class WifiTrafficPoller {
     /* Tracks last reported data activity */
     private int mDataActivity;
 
-    private final List<Messenger> mClients = new ArrayList<Messenger>();
+    private final HashMap<Integer, ITrafficStateCallback> mRegisteredCallbacks = new HashMap<>();
     // err on the side of updating at boot since screen on broadcast may be missed
     // the first time
     private AtomicBoolean mScreenOn = new AtomicBoolean(true);
@@ -105,14 +105,29 @@ public class WifiTrafficPoller {
                 }, filter);
     }
 
-    /** */
-    public void addClient(Messenger client) {
-        Message.obtain(mTrafficHandler, ADD_CLIENT, client).sendToTarget();
+    /**
+     * Add a new callback to the traffic poller.
+     */
+    public void addCallback(ITrafficStateCallback callback, int callbackIdentifier) {
+        mRegisteredCallbacks.put(callbackIdentifier, callback);
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Adding callback. Num callbacks: " + mRegisteredCallbacks.size());
+        }
+        if (mRegisteredCallbacks.size() > NUM_CALLBACKS_WTF_LIMIT) {
+            Log.wtf(TAG, "Too many traffic poller callbacks: " + mRegisteredCallbacks.size());
+        } else if (mRegisteredCallbacks.size() > NUM_CALLBACKS_WARN_LIMIT) {
+            Log.w(TAG, "Too many traffic poller callbacks: " + mRegisteredCallbacks.size());
+        }
     }
 
-    /** */
-    public void removeClient(Messenger client) {
-        Message.obtain(mTrafficHandler, REMOVE_CLIENT, client).sendToTarget();
+    /**
+     * Remove an existing callback to the traffic poller.
+     */
+    public void removeCallback(int callbackIdentifier) {
+        mRegisteredCallbacks.remove(callbackIdentifier);
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Removing callback. Num callbacks: " + mRegisteredCallbacks.size());
+        }
     }
 
     void enableVerboseLogging(int verbose) {
@@ -151,7 +166,7 @@ public class WifiTrafficPoller {
                         Log.d(TAG, "TRAFFIC_STATS_POLL "
                                 + mEnableTrafficStatsPoll + " Token "
                                 + Integer.toString(mTrafficStatsPollToken)
-                                + " num clients " + mClients.size());
+                                + " num clients " + mRegisteredCallbacks.size());
                     }
                     if (msg.arg1 == mTrafficStatsPollToken) {
                         ifaceName = mWifiNative.getClientInterfaceName();
@@ -162,18 +177,7 @@ public class WifiTrafficPoller {
                         }
                     }
                     break;
-                case ADD_CLIENT:
-                    mClients.add((Messenger) msg.obj);
-                    if (mVerboseLoggingEnabled) {
-                        Log.d(TAG, "ADD_CLIENT: "
-                                + Integer.toString(mClients.size()));
-                    }
-                    break;
-                case REMOVE_CLIENT:
-                    mClients.remove(msg.obj);
-                    break;
             }
-
         }
     }
 
@@ -193,8 +197,9 @@ public class WifiTrafficPoller {
     private void notifyOnDataActivity(@NonNull String ifaceName) {
         long sent, received;
         long preTxPkts = mTxPkts, preRxPkts = mRxPkts;
-        int dataActivity = WifiManager.DATA_ACTIVITY_NONE;
+        int dataActivity = WifiManager.TrafficStateCallback.DATA_ACTIVITY_NONE;
 
+        // TODO (b/111691443): Use WifiInfo instead of making the native calls here.
         mTxPkts = mWifiNative.getTxPackets(ifaceName);
         mRxPkts = mWifiNative.getRxPackets(ifaceName);
 
@@ -209,10 +214,10 @@ public class WifiTrafficPoller {
             sent = mTxPkts - preTxPkts;
             received = mRxPkts - preRxPkts;
             if (sent > 0) {
-                dataActivity |= WifiManager.DATA_ACTIVITY_OUT;
+                dataActivity |= WifiManager.TrafficStateCallback.DATA_ACTIVITY_OUT;
             }
             if (received > 0) {
-                dataActivity |= WifiManager.DATA_ACTIVITY_IN;
+                dataActivity |= WifiManager.TrafficStateCallback.DATA_ACTIVITY_IN;
             }
 
             if (dataActivity != mDataActivity && mScreenOn.get()) {
@@ -221,12 +226,9 @@ public class WifiTrafficPoller {
                     Log.e(TAG, "notifying of data activity "
                             + Integer.toString(mDataActivity));
                 }
-                for (Messenger client : mClients) {
-                    Message msg = Message.obtain();
-                    msg.what = WifiManager.DATA_ACTIVITY_NOTIFICATION;
-                    msg.arg1 = mDataActivity;
+                for (ITrafficStateCallback callback : mRegisteredCallbacks.values()) {
                     try {
-                        client.send(msg);
+                        callback.onStateChanged(mDataActivity);
                     } catch (RemoteException e) {
                         // Failed to reach, skip
                         // Client removal is handled in WifiService
@@ -242,6 +244,7 @@ public class WifiTrafficPoller {
         pw.println("mTxPkts " + mTxPkts);
         pw.println("mRxPkts " + mRxPkts);
         pw.println("mDataActivity " + mDataActivity);
+        pw.println("mRegisteredCallbacks " + mRegisteredCallbacks);
     }
 
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
index addaf354b..f4a97c2b0 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiServiceImplTest.java
@@ -66,6 +66,7 @@ import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.net.Uri;
 import android.net.wifi.ISoftApCallback;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.ScanResult;
 import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiConfiguration.KeyMgmt;
@@ -137,6 +138,7 @@ public class WifiServiceImplTest {
     private static final int TEST_UID = 1200000;
     private static final int OTHER_TEST_UID = 1300000;
     private static final int TEST_USER_HANDLE = 13;
+    private static final int TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER = 17;
     private static final String WIFI_IFACE_NAME = "wlan0";
     private static final String TEST_COUNTRY_CODE = "US";
 
@@ -197,6 +199,7 @@ public class WifiServiceImplTest {
     @Mock PowerProfile mPowerProfile;
     @Mock WifiTrafficPoller mWifiTrafficPolller;
     @Mock ScanRequestProxy mScanRequestProxy;
+    @Mock ITrafficStateCallback mTrafficStateCallback;
 
     @Spy FakeWifiLog mLog;
 
@@ -2865,4 +2868,101 @@ public class WifiServiceImplTest {
             fail("Expected Security exception");
         } catch (SecurityException e) { }
     }
+
+    /**
+     * Verify that a call to registerTrafficStateCallback throws a SecurityException if the caller
+     * does not have NETWORK_SETTINGS permission.
+     */
+    @Test
+    public void registerTrafficStateCallbackThrowsSecurityExceptionOnMissingPermissions() {
+        doThrow(new SecurityException()).when(mContext)
+                .enforceCallingOrSelfPermission(eq(android.Manifest.permission.NETWORK_SETTINGS),
+                        eq("WifiService"));
+        try {
+            mWifiServiceImpl.registerTrafficStateCallback(mAppBinder, mTrafficStateCallback,
+                    TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    /**
+     * Verify that a call to registerTrafficStateCallback throws an IllegalArgumentException if the
+     * parameters are not provided.
+     */
+    @Test
+    public void registerTrafficStateCallbackThrowsIllegalArgumentExceptionOnInvalidArguments() {
+        try {
+            mWifiServiceImpl.registerTrafficStateCallback(
+                    mAppBinder, null, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected IllegalArgumentException");
+        } catch (IllegalArgumentException expected) {
+        }
+    }
+
+    /**
+     * Verify that a call to unregisterTrafficStateCallback throws a SecurityException if the caller
+     * does not have NETWORK_SETTINGS permission.
+     */
+    @Test
+    public void unregisterTrafficStateCallbackThrowsSecurityExceptionOnMissingPermissions() {
+        doThrow(new SecurityException()).when(mContext)
+                .enforceCallingOrSelfPermission(eq(android.Manifest.permission.NETWORK_SETTINGS),
+                        eq("WifiService"));
+        try {
+            mWifiServiceImpl.unregisterTrafficStateCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+            fail("expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    /**
+     * Verify that registerTrafficStateCallback adds callback to {@link WifiTrafficPoller}.
+     */
+    @Test
+    public void registerTrafficStateCallbackAndVerify() throws Exception {
+        mWifiServiceImpl.registerTrafficStateCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        mLooper.dispatchAll();
+        verify(mWifiTrafficPoller).addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+    }
+
+    /**
+     * Verify that unregisterTrafficStateCallback removes callback from {@link WifiTrafficPoller}.
+     */
+    @Test
+    public void unregisterTrafficStateCallbackAndVerify() throws Exception {
+        mWifiServiceImpl.unregisterTrafficStateCallback(0);
+        mLooper.dispatchAll();
+        verify(mWifiTrafficPoller).removeCallback(0);
+    }
+
+    /**
+     * Verify that wifi service registers for callers BinderDeath event
+     */
+    @Test
+    public void registersForBinderDeathOnRegisterTrafficStateCallback() throws Exception {
+        mWifiServiceImpl.registerTrafficStateCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        mLooper.dispatchAll();
+        verify(mAppBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+    }
+
+    /**
+     * Verify that we remove the traffic state callback on receiving BinderDied event.
+     */
+    @Test
+    public void unregistersTrafficStateCallbackOnBinderDied() throws Exception {
+        ArgumentCaptor<IBinder.DeathRecipient> drCaptor =
+                ArgumentCaptor.forClass(IBinder.DeathRecipient.class);
+        mWifiServiceImpl.registerTrafficStateCallback(
+                mAppBinder, mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        verify(mAppBinder).linkToDeath(drCaptor.capture(), anyInt());
+
+        drCaptor.getValue().binderDied();
+        mLooper.dispatchAll();
+        verify(mAppBinder).unlinkToDeath(drCaptor.getValue(), 0);
+        verify(mWifiTrafficPoller).removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java b/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
index 40281f8ba..b00c4f74c 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiTrafficPollerTest.java
@@ -15,32 +15,31 @@
  */
 package com.android.server.wifi;
 
-import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import org.junit.Before;
-import org.junit.Test;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
 import android.net.NetworkInfo;
+import android.net.wifi.ITrafficStateCallback;
 import android.net.wifi.WifiManager;
-import android.os.Handler;
 import android.os.Message;
-import android.os.Messenger;
+import android.os.RemoteException;
 import android.os.test.TestLooper;
 import android.support.test.filters.SmallTest;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
 /**
  * Unit tests for {@link com.android.server.wifi.WifiTrafficPoller}.
  */
@@ -49,15 +48,14 @@ public class WifiTrafficPollerTest {
     public static final String TAG = "WifiTrafficPollerTest";
 
     private TestLooper mLooper;
-    private Handler mHandler;
     private WifiTrafficPoller mWifiTrafficPoller;
     private BroadcastReceiver mReceiver;
     private Intent mIntent;
-    private Messenger mMessenger;
     private final static String IFNAME = "wlan0";
     private final static long DEFAULT_PACKET_COUNT = 10;
     private final static long TX_PACKET_COUNT = 40;
     private final static long RX_PACKET_COUNT = 50;
+    private static final int TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER = 14;
 
     final ArgumentCaptor<Message> mMessageCaptor = ArgumentCaptor.forClass(Message.class);
     final ArgumentCaptor<BroadcastReceiver> mBroadcastReceiverCaptor =
@@ -66,6 +64,7 @@ public class WifiTrafficPollerTest {
     @Mock Context mContext;
     @Mock WifiNative mWifiNative;
     @Mock NetworkInfo mNetworkInfo;
+    @Mock ITrafficStateCallback mTrafficStateCallback;
 
     /**
      * Called before each test
@@ -74,8 +73,6 @@ public class WifiTrafficPollerTest {
     public void setUp() throws Exception {
         // Ensure looper exists
         mLooper = new TestLooper();
-        mHandler = spy(new Handler(mLooper.getLooper()));
-        mMessenger = new Messenger(mHandler);
         MockitoAnnotations.initMocks(this);
 
         when(mWifiNative.getTxPackets(any(String.class))).thenReturn(DEFAULT_PACKET_COUNT,
@@ -98,12 +95,6 @@ public class WifiTrafficPollerTest {
                 NetworkInfo.DetailedState.CONNECTED);
     }
 
-    private void registerClient() {
-        // Register Client to verify that Tx/RX packet message is properly received.
-        mWifiTrafficPoller.addClient(mMessenger);
-        mLooper.dispatchAll();
-    }
-
     private void triggerForUpdatedInformationOfData(String actionScreen,
             NetworkInfo.DetailedState networkState) {
         when(mNetworkInfo.getDetailedState()).thenReturn(NetworkInfo.DetailedState.DISCONNECTED);
@@ -122,13 +113,15 @@ public class WifiTrafficPollerTest {
      * Verify that StartTrafficStatsPolling should not happen in case a network is not connected
      */
     @Test
-    public void testNotStartTrafficStatsPollingWithDisconnected() {
-        registerClient();
+    public void testNotStartTrafficStatsPollingWithDisconnected() throws RemoteException {
+        // Register Client to verify that Tx/RX packet message is properly received.
+        mWifiTrafficPoller.addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.DISCONNECTED);
 
         // Client should not get any message when the network is disconnected
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
     }
 
     /**
@@ -136,22 +129,26 @@ public class WifiTrafficPollerTest {
      * available.
      */
     @Test
-    public void testStartTrafficStatsPollingWithScreenOn() {
-        registerClient();
+    public void testStartTrafficStatsPollingWithScreenOn() throws RemoteException {
+        // Register Client to verify that Tx/RX packet message is properly received.
+        mWifiTrafficPoller.addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
 
         // Client should get the DATA_ACTIVITY_NOTIFICATION
-        verify(mHandler).handleMessage(mMessageCaptor.capture());
-        assertEquals(WifiManager.DATA_ACTIVITY_NOTIFICATION, mMessageCaptor.getValue().what);
+        verify(mTrafficStateCallback).onStateChanged(
+                WifiManager.TrafficStateCallback.DATA_ACTIVITY_INOUT);
     }
 
     /**
      * Verify that StartTrafficStatsPolling should not happen in case screen is off.
      */
     @Test
-    public void testNotStartTrafficStatsPollingWithScreenOff() {
-        registerClient();
+    public void testNotStartTrafficStatsPollingWithScreenOff() throws RemoteException {
+        // Register Client to verify that Tx/RX packet message is properly received.
+        mWifiTrafficPoller.addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_OFF,
                 NetworkInfo.DetailedState.CONNECTED);
 
@@ -159,22 +156,43 @@ public class WifiTrafficPollerTest {
         mLooper.dispatchAll();
 
         // Client should not get any message when the screen is off
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
     }
 
     /**
-     * Verify that remove client message should be handled
+     * Verify that remove client should be handled
      */
     @Test
-    public void testRemoveClient() {
-        registerClient();
-        mWifiTrafficPoller.removeClient(mMessenger);
+    public void testRemoveClient() throws RemoteException {
+        // Register Client to verify that Tx/RX packet message is properly received.
+        mWifiTrafficPoller.addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
         mLooper.dispatchAll();
 
         triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
                 NetworkInfo.DetailedState.CONNECTED);
 
         // Client should not get any message after the client is removed.
-        verify(mHandler, never()).handleMessage(any(Message.class));
+        verify(mTrafficStateCallback, never()).onStateChanged(anyInt());
+    }
+
+    /**
+     * Verify that remove client ignores when callback identifier is wrong.
+     */
+    @Test
+    public void testRemoveClientWithWrongIdentifier() throws RemoteException {
+        // Register Client to verify that Tx/RX packet message is properly received.
+        mWifiTrafficPoller.addCallback(
+                mTrafficStateCallback, TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER);
+        mWifiTrafficPoller.removeCallback(TEST_TRAFFIC_STATE_CALLBACK_IDENTIFIER + 5);
+        mLooper.dispatchAll();
+
+        triggerForUpdatedInformationOfData(Intent.ACTION_SCREEN_ON,
+                NetworkInfo.DetailedState.CONNECTED);
+
+        // Client should get the DATA_ACTIVITY_NOTIFICATION
+        verify(mTrafficStateCallback).onStateChanged(
+                WifiManager.TrafficStateCallback.DATA_ACTIVITY_INOUT);
     }
 }