Merge "Add nullable parameter to readEmbeddedBuffer."
diff --git a/core/jni/android_os_HwBinder.h b/core/jni/android_os_HwBinder.h
index 2ebc381..fa8fe01 100644
--- a/core/jni/android_os_HwBinder.h
+++ b/core/jni/android_os_HwBinder.h
@@ -24,7 +24,7 @@
 
 namespace android {
 
-struct JHwBinder : public hardware::BBinder {
+struct JHwBinder : public hardware::BHwBinder {
     static void InitClass(JNIEnv *env);
 
     static sp<JHwBinder> SetNativeContext(
diff --git a/services/core/java/com/android/server/connectivity/Tethering.java b/services/core/java/com/android/server/connectivity/Tethering.java
index 57381f2..0c80166 100644
--- a/services/core/java/com/android/server/connectivity/Tethering.java
+++ b/services/core/java/com/android/server/connectivity/Tethering.java
@@ -32,7 +32,6 @@
 import android.content.res.Resources;
 import android.hardware.usb.UsbManager;
 import android.net.ConnectivityManager;
-import android.net.ConnectivityManager.NetworkCallback;
 import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
 import android.net.LinkProperties;
@@ -73,6 +72,7 @@
 import com.android.server.connectivity.tethering.IPv6TetheringCoordinator;
 import com.android.server.connectivity.tethering.IPv6TetheringInterfaceServices;
 import com.android.server.connectivity.tethering.TetherInterfaceStateMachine;
+import com.android.server.connectivity.tethering.UpstreamNetworkMonitor;
 import com.android.server.net.BaseNetworkObserver;
 
 import java.io.FileDescriptor;
@@ -1031,249 +1031,6 @@
         }
     }
 
-    /**
-     * A class to centralize all the network and link properties information
-     * pertaining to the current and any potential upstream network.
-     *
-     * Calling #start() registers two callbacks: one to track the system default
-     * network and a second to specifically observe TYPE_MOBILE_DUN networks.
-     *
-     * The methods and data members of this class are only to be accessed and
-     * modified from the tethering master state machine thread. Any other
-     * access semantics would necessitate the addition of locking.
-     *
-     * TODO: Investigate whether more "upstream-specific" logic/functionality
-     * could/should be moved here.
-     */
-    public class UpstreamNetworkMonitor {
-        public static final int EVENT_ON_AVAILABLE      = 1;
-        public static final int EVENT_ON_CAPABILITIES   = 2;
-        public static final int EVENT_ON_LINKPROPERTIES = 3;
-        public static final int EVENT_ON_LOST           = 4;
-
-        private final Context mContext;
-        private final StateMachine mTarget;
-        private final int mWhat;
-        private final HashMap<Network, NetworkState> mNetworkMap = new HashMap<>();
-        private ConnectivityManager mCM;
-        private NetworkCallback mDefaultNetworkCallback;
-        private NetworkCallback mDunTetheringCallback;
-        private NetworkCallback mMobileNetworkCallback;
-        private boolean mDunRequired;
-
-        public UpstreamNetworkMonitor(Context ctx, StateMachine tgt, int what) {
-            mContext = ctx;
-            mTarget = tgt;
-            mWhat = what;
-        }
-
-        public void start() {
-            stop();
-
-            mDefaultNetworkCallback = new UpstreamNetworkCallback();
-            cm().registerDefaultNetworkCallback(mDefaultNetworkCallback);
-
-            final NetworkRequest dunTetheringRequest = new NetworkRequest.Builder()
-                    .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
-                    .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
-                    .addCapability(NetworkCapabilities.NET_CAPABILITY_DUN)
-                    .build();
-            mDunTetheringCallback = new UpstreamNetworkCallback();
-            cm().registerNetworkCallback(dunTetheringRequest, mDunTetheringCallback);
-        }
-
-        public void stop() {
-            releaseMobileNetworkRequest();
-
-            releaseCallback(mDefaultNetworkCallback);
-            mDefaultNetworkCallback = null;
-
-            releaseCallback(mDunTetheringCallback);
-            mDunTetheringCallback = null;
-
-            mNetworkMap.clear();
-        }
-
-        public void mobileUpstreamRequiresDun(boolean dunRequired) {
-            final boolean valueChanged = (mDunRequired != dunRequired);
-            mDunRequired = dunRequired;
-            if (valueChanged && mobileNetworkRequested()) {
-                releaseMobileNetworkRequest();
-                registerMobileNetworkRequest();
-            }
-        }
-
-        public boolean mobileNetworkRequested() {
-            return (mMobileNetworkCallback != null);
-        }
-
-        public void registerMobileNetworkRequest() {
-            if (mMobileNetworkCallback != null) return;
-
-            final NetworkRequest.Builder builder = new NetworkRequest.Builder()
-                    .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-            if (mDunRequired) {
-                builder.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
-                       .addCapability(NetworkCapabilities.NET_CAPABILITY_DUN);
-            } else {
-                builder.addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET);
-            }
-            final NetworkRequest mobileUpstreamRequest = builder.build();
-
-            // The existing default network and DUN callbacks will be notified.
-            // Therefore, to avoid duplicate notifications, we only register a no-op.
-            mMobileNetworkCallback = new NetworkCallback();
-
-            // TODO: Change the timeout from 0 (no onUnavailable callback) to use some
-            // moderate callback time (once timeout callbacks are implemented). This might
-            // be useful for updating some UI. Additionally, we should definitely log a
-            // message to aid in any subsequent debugging
-            if (DBG) Log.d(TAG, "requesting mobile upstream network: " + mobileUpstreamRequest);
-
-            cm().requestNetwork(mobileUpstreamRequest, mMobileNetworkCallback);
-        }
-
-        public void releaseMobileNetworkRequest() {
-            if (mMobileNetworkCallback == null) return;
-
-            cm().unregisterNetworkCallback(mMobileNetworkCallback);
-            mMobileNetworkCallback = null;
-        }
-
-        public NetworkState lookup(Network network) {
-            return (network != null) ? mNetworkMap.get(network) : null;
-        }
-
-        private void handleAvailable(Network network) {
-            if (VDBG) {
-                Log.d(TAG, "EVENT_ON_AVAILABLE for " + network);
-            }
-            if (!mNetworkMap.containsKey(network)) {
-                mNetworkMap.put(network,
-                        new NetworkState(null, null, null, network, null, null));
-            }
-
-            final ConnectivityManager cm = cm();
-
-            if (mDefaultNetworkCallback != null) {
-                cm.requestNetworkCapabilities(mDefaultNetworkCallback);
-                cm.requestLinkProperties(mDefaultNetworkCallback);
-            }
-
-            // Requesting updates for mDunTetheringCallback is not
-            // necessary. Because it's a listen, it will already have
-            // heard all NetworkCapabilities and LinkProperties updates
-            // since UpstreamNetworkMonitor was started. Because we
-            // start UpstreamNetworkMonitor before chooseUpstreamType()
-            // is ever invoked (it can register a DUN request) this is
-            // mostly safe. However, if a DUN network is already up for
-            // some reason (unlikely, because DUN is restricted and,
-            // unless the DUN network is shared with another APN, only
-            // the system can request it and this is the only part of
-            // the system that requests it) we won't know its
-            // LinkProperties or NetworkCapabilities.
-
-            notifyTarget(EVENT_ON_AVAILABLE, network);
-        }
-
-        private void handleNetCap(Network network, NetworkCapabilities newNc) {
-            if (!mNetworkMap.containsKey(network)) {
-                // Ignore updates for networks for which we have not yet
-                // received onAvailable() - which should never happen -
-                // or for which we have already received onLost().
-                return;
-            }
-            if (VDBG) {
-                Log.d(TAG, String.format("EVENT_ON_CAPABILITIES for %s: %s",
-                        network, newNc));
-            }
-
-            final NetworkState prev = mNetworkMap.get(network);
-            mNetworkMap.put(network,
-                    new NetworkState(null, prev.linkProperties, newNc,
-                                     network, null, null));
-            notifyTarget(EVENT_ON_CAPABILITIES, network);
-        }
-
-        private void handleLinkProp(Network network, LinkProperties newLp) {
-            if (!mNetworkMap.containsKey(network)) {
-                // Ignore updates for networks for which we have not yet
-                // received onAvailable() - which should never happen -
-                // or for which we have already received onLost().
-                return;
-            }
-            if (VDBG) {
-                Log.d(TAG, String.format("EVENT_ON_LINKPROPERTIES for %s: %s",
-                        network, newLp));
-            }
-
-            final NetworkState prev = mNetworkMap.get(network);
-            mNetworkMap.put(network,
-                    new NetworkState(null, newLp, prev.networkCapabilities,
-                                     network, null, null));
-            notifyTarget(EVENT_ON_LINKPROPERTIES, network);
-        }
-
-        private void handleLost(Network network) {
-            if (!mNetworkMap.containsKey(network)) {
-                // Ignore updates for networks for which we have not yet
-                // received onAvailable() - which should never happen -
-                // or for which we have already received onLost().
-                return;
-            }
-            if (VDBG) {
-                Log.d(TAG, "EVENT_ON_LOST for " + network);
-            }
-            notifyTarget(EVENT_ON_LOST, mNetworkMap.remove(network));
-        }
-
-        // Fetch (and cache) a ConnectivityManager only if and when we need one.
-        private ConnectivityManager cm() {
-            if (mCM == null) {
-                mCM = mContext.getSystemService(ConnectivityManager.class);
-            }
-            return mCM;
-        }
-
-        /**
-         * A NetworkCallback class that relays information of interest to the
-         * tethering master state machine thread for subsequent processing.
-         */
-        private class UpstreamNetworkCallback extends NetworkCallback {
-            @Override
-            public void onAvailable(Network network) {
-                mTarget.getHandler().post(() -> handleAvailable(network));
-            }
-
-            @Override
-            public void onCapabilitiesChanged(Network network, NetworkCapabilities newNc) {
-                mTarget.getHandler().post(() -> handleNetCap(network, newNc));
-            }
-
-            @Override
-            public void onLinkPropertiesChanged(Network network, LinkProperties newLp) {
-                mTarget.getHandler().post(() -> handleLinkProp(network, newLp));
-            }
-
-            @Override
-            public void onLost(Network network) {
-                mTarget.getHandler().post(() -> handleLost(network));
-            }
-        }
-
-        private void releaseCallback(NetworkCallback cb) {
-            if (cb != null) cm().unregisterNetworkCallback(cb);
-        }
-
-        private void notifyTarget(int which, Network network) {
-            notifyTarget(which, mNetworkMap.get(network));
-        }
-
-        private void notifyTarget(int which, NetworkState netstate) {
-            mTarget.sendMessage(mWhat, which, 0, netstate);
-        }
-    }
-
     // Needed because the canonical source of upstream truth is just the
     // upstream interface name, |mCurrentUpstreamIface|.  This is ripe for
     // future simplification, once the upstream Network is canonical.
diff --git a/services/core/java/com/android/server/connectivity/tethering/UpstreamNetworkMonitor.java b/services/core/java/com/android/server/connectivity/tethering/UpstreamNetworkMonitor.java
new file mode 100644
index 0000000..eb79211
--- /dev/null
+++ b/services/core/java/com/android/server/connectivity/tethering/UpstreamNetworkMonitor.java
@@ -0,0 +1,280 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity.tethering;
+
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.net.ConnectivityManager.NetworkCallback;
+import android.net.LinkProperties;
+import android.net.Network;
+import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
+import android.net.NetworkState;
+import android.util.Log;
+
+import com.android.internal.util.StateMachine;
+
+import java.util.HashMap;
+
+
+/**
+ * A class to centralize all the network and link properties information
+ * pertaining to the current and any potential upstream network.
+ *
+ * Calling #start() registers two callbacks: one to track the system default
+ * network and a second to specifically observe TYPE_MOBILE_DUN networks.
+ *
+ * The methods and data members of this class are only to be accessed and
+ * modified from the tethering master state machine thread. Any other
+ * access semantics would necessitate the addition of locking.
+ *
+ * TODO: Move upstream selection logic here.
+ *
+ * @hide
+ */
+public class UpstreamNetworkMonitor {
+    private static final String TAG = UpstreamNetworkMonitor.class.getSimpleName();
+    private static final boolean DBG = false;
+    private static final boolean VDBG = false;
+
+    public static final int EVENT_ON_AVAILABLE      = 1;
+    public static final int EVENT_ON_CAPABILITIES   = 2;
+    public static final int EVENT_ON_LINKPROPERTIES = 3;
+    public static final int EVENT_ON_LOST           = 4;
+
+    private final Context mContext;
+    private final StateMachine mTarget;
+    private final int mWhat;
+    private final HashMap<Network, NetworkState> mNetworkMap = new HashMap<>();
+    private ConnectivityManager mCM;
+    private NetworkCallback mDefaultNetworkCallback;
+    private NetworkCallback mDunTetheringCallback;
+    private NetworkCallback mMobileNetworkCallback;
+    private boolean mDunRequired;
+
+    public UpstreamNetworkMonitor(Context ctx, StateMachine tgt, int what) {
+        mContext = ctx;
+        mTarget = tgt;
+        mWhat = what;
+    }
+
+    public void start() {
+        stop();
+
+        mDefaultNetworkCallback = new UpstreamNetworkCallback();
+        cm().registerDefaultNetworkCallback(mDefaultNetworkCallback);
+
+        final NetworkRequest dunTetheringRequest = new NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_DUN)
+                .build();
+        mDunTetheringCallback = new UpstreamNetworkCallback();
+        cm().registerNetworkCallback(dunTetheringRequest, mDunTetheringCallback);
+    }
+
+    public void stop() {
+        releaseMobileNetworkRequest();
+
+        releaseCallback(mDefaultNetworkCallback);
+        mDefaultNetworkCallback = null;
+
+        releaseCallback(mDunTetheringCallback);
+        mDunTetheringCallback = null;
+
+        mNetworkMap.clear();
+    }
+
+    public void mobileUpstreamRequiresDun(boolean dunRequired) {
+        final boolean valueChanged = (mDunRequired != dunRequired);
+        mDunRequired = dunRequired;
+        if (valueChanged && mobileNetworkRequested()) {
+            releaseMobileNetworkRequest();
+            registerMobileNetworkRequest();
+        }
+    }
+
+    public boolean mobileNetworkRequested() {
+        return (mMobileNetworkCallback != null);
+    }
+
+    public void registerMobileNetworkRequest() {
+        if (mMobileNetworkCallback != null) return;
+
+        final NetworkRequest.Builder builder = new NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
+        if (mDunRequired) {
+            builder.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                   .addCapability(NetworkCapabilities.NET_CAPABILITY_DUN);
+        } else {
+            builder.addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET);
+        }
+        final NetworkRequest mobileUpstreamRequest = builder.build();
+
+        // The existing default network and DUN callbacks will be notified.
+        // Therefore, to avoid duplicate notifications, we only register a no-op.
+        mMobileNetworkCallback = new NetworkCallback();
+
+        // TODO: Change the timeout from 0 (no onUnavailable callback) to use some
+        // moderate callback time (once timeout callbacks are implemented). This might
+        // be useful for updating some UI. Additionally, we should definitely log a
+        // message to aid in any subsequent debugging
+        if (DBG) Log.d(TAG, "requesting mobile upstream network: " + mobileUpstreamRequest);
+
+        cm().requestNetwork(mobileUpstreamRequest, mMobileNetworkCallback);
+    }
+
+    public void releaseMobileNetworkRequest() {
+        if (mMobileNetworkCallback == null) return;
+
+        cm().unregisterNetworkCallback(mMobileNetworkCallback);
+        mMobileNetworkCallback = null;
+    }
+
+    public NetworkState lookup(Network network) {
+        return (network != null) ? mNetworkMap.get(network) : null;
+    }
+
+    private void handleAvailable(Network network) {
+        if (VDBG) {
+            Log.d(TAG, "EVENT_ON_AVAILABLE for " + network);
+        }
+        if (!mNetworkMap.containsKey(network)) {
+            mNetworkMap.put(network,
+                    new NetworkState(null, null, null, network, null, null));
+        }
+
+        final ConnectivityManager cm = cm();
+
+        if (mDefaultNetworkCallback != null) {
+            cm.requestNetworkCapabilities(mDefaultNetworkCallback);
+            cm.requestLinkProperties(mDefaultNetworkCallback);
+        }
+
+        // Requesting updates for mDunTetheringCallback is not
+        // necessary. Because it's a listen, it will already have
+        // heard all NetworkCapabilities and LinkProperties updates
+        // since UpstreamNetworkMonitor was started. Because we
+        // start UpstreamNetworkMonitor before chooseUpstreamType()
+        // is ever invoked (it can register a DUN request) this is
+        // mostly safe. However, if a DUN network is already up for
+        // some reason (unlikely, because DUN is restricted and,
+        // unless the DUN network is shared with another APN, only
+        // the system can request it and this is the only part of
+        // the system that requests it) we won't know its
+        // LinkProperties or NetworkCapabilities.
+
+        notifyTarget(EVENT_ON_AVAILABLE, network);
+    }
+
+    private void handleNetCap(Network network, NetworkCapabilities newNc) {
+        if (!mNetworkMap.containsKey(network)) {
+            // Ignore updates for networks for which we have not yet
+            // received onAvailable() - which should never happen -
+            // or for which we have already received onLost().
+            return;
+        }
+        if (VDBG) {
+            Log.d(TAG, String.format("EVENT_ON_CAPABILITIES for %s: %s",
+                    network, newNc));
+        }
+
+        final NetworkState prev = mNetworkMap.get(network);
+        mNetworkMap.put(network,
+                new NetworkState(null, prev.linkProperties, newNc,
+                                 network, null, null));
+        notifyTarget(EVENT_ON_CAPABILITIES, network);
+    }
+
+    private void handleLinkProp(Network network, LinkProperties newLp) {
+        if (!mNetworkMap.containsKey(network)) {
+            // Ignore updates for networks for which we have not yet
+            // received onAvailable() - which should never happen -
+            // or for which we have already received onLost().
+            return;
+        }
+        if (VDBG) {
+            Log.d(TAG, String.format("EVENT_ON_LINKPROPERTIES for %s: %s",
+                    network, newLp));
+        }
+
+        final NetworkState prev = mNetworkMap.get(network);
+        mNetworkMap.put(network,
+                new NetworkState(null, newLp, prev.networkCapabilities,
+                                 network, null, null));
+        notifyTarget(EVENT_ON_LINKPROPERTIES, network);
+    }
+
+    private void handleLost(Network network) {
+        if (!mNetworkMap.containsKey(network)) {
+            // Ignore updates for networks for which we have not yet
+            // received onAvailable() - which should never happen -
+            // or for which we have already received onLost().
+            return;
+        }
+        if (VDBG) {
+            Log.d(TAG, "EVENT_ON_LOST for " + network);
+        }
+        notifyTarget(EVENT_ON_LOST, mNetworkMap.remove(network));
+    }
+
+    // Fetch (and cache) a ConnectivityManager only if and when we need one.
+    private ConnectivityManager cm() {
+        if (mCM == null) {
+            mCM = mContext.getSystemService(ConnectivityManager.class);
+        }
+        return mCM;
+    }
+
+    /**
+     * A NetworkCallback class that relays information of interest to the
+     * tethering master state machine thread for subsequent processing.
+     */
+    private class UpstreamNetworkCallback extends NetworkCallback {
+        @Override
+        public void onAvailable(Network network) {
+            mTarget.getHandler().post(() -> handleAvailable(network));
+        }
+
+        @Override
+        public void onCapabilitiesChanged(Network network, NetworkCapabilities newNc) {
+            mTarget.getHandler().post(() -> handleNetCap(network, newNc));
+        }
+
+        @Override
+        public void onLinkPropertiesChanged(Network network, LinkProperties newLp) {
+            mTarget.getHandler().post(() -> handleLinkProp(network, newLp));
+        }
+
+        @Override
+        public void onLost(Network network) {
+            mTarget.getHandler().post(() -> handleLost(network));
+        }
+    }
+
+    private void releaseCallback(NetworkCallback cb) {
+        if (cb != null) cm().unregisterNetworkCallback(cb);
+    }
+
+    private void notifyTarget(int which, Network network) {
+        notifyTarget(which, mNetworkMap.get(network));
+    }
+
+    private void notifyTarget(int which, NetworkState netstate) {
+        mTarget.sendMessage(mWhat, which, 0, netstate);
+    }
+}
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index d62c30d..46b6403 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -1103,6 +1103,29 @@
         UNAVAILABLE
     }
 
+    private static class CallbackInfo {
+        public final CallbackState state;
+        public final Network network;
+        public final Object arg;
+        public CallbackInfo(CallbackState s, Network n, Object o) {
+            state = s; network = n; arg = o;
+        }
+        public String toString() {
+            return String.format("%s (%s)", state, network);
+        }
+        @Override
+        public boolean equals(Object o) {
+            if (!(o instanceof CallbackInfo)) return false;
+            // Ignore timeMs, since it's unpredictable.
+            CallbackInfo other = (CallbackInfo) o;
+            return (state == other.state) && Objects.equals(network, other.network);
+        }
+        @Override
+        public int hashCode() {
+            return Objects.hash(state, network);
+        }
+    }
+
     /**
      * Utility NetworkCallback for testing. The caller must explicitly test for all the callbacks
      * this class receives, by calling expectCallback() exactly once each time a callback is
@@ -1114,21 +1137,6 @@
         // the linger timeout.
         private final static int TIMEOUT_MS = 50;
 
-        private class CallbackInfo {
-            public final CallbackState state;
-            public final Network network;
-            public Object arg;
-            public CallbackInfo(CallbackState s, Network n, Object o) {
-                state = s; network = n; arg = o;
-            }
-            public String toString() { return String.format("%s (%s)", state, network); }
-            public boolean equals(Object o) {
-                if (!(o instanceof CallbackInfo)) return false;
-                // Ignore timeMs, since it's unpredictable.
-                CallbackInfo other = (CallbackInfo) o;
-                return state == other.state && Objects.equals(network, other.network);
-            }
-        }
         private final LinkedBlockingQueue<CallbackInfo> mCallbacks = new LinkedBlockingQueue<>();
 
         protected void setLastCallback(CallbackState state, Network network, Object o) {
@@ -1155,17 +1163,24 @@
             setLastCallback(CallbackState.LOST, network, null);
         }
 
+        CallbackInfo nextCallback(int timeoutMs) {
+            CallbackInfo cb = null;
+            try {
+                cb = mCallbacks.poll(timeoutMs, TimeUnit.MILLISECONDS);
+            } catch (InterruptedException e) {
+            }
+            if (cb == null) {
+                // LinkedBlockingQueue.poll() returns null if it timeouts.
+                fail("Did not receive callback after " + timeoutMs + "ms");
+            }
+            return cb;
+        }
+
         void expectCallback(CallbackState state, MockNetworkAgent mockAgent, int timeoutMs) {
             CallbackInfo expected = new CallbackInfo(
                     state, (mockAgent != null) ? mockAgent.getNetwork() : null, 0);
-            CallbackInfo actual;
-            try {
-                actual = mCallbacks.poll(timeoutMs, TimeUnit.MILLISECONDS);
-                assertEquals("Unexpected callback:", expected, actual);
-            } catch (InterruptedException e) {
-                fail("Did not receive expected " + expected + " after " + TIMEOUT_MS + "ms");
-                actual = null;  // Or the compiler can't tell it's never used uninitialized.
-            }
+            CallbackInfo actual = nextCallback(timeoutMs);
+            assertEquals("Unexpected callback:", expected, actual);
             if (state == CallbackState.LOSING) {
                 String msg = String.format(
                         "Invalid linger time value %d, must be between %d and %d",