Merge changes from topic 'packet-keepalive-fixes' into mnc-dr-dev

* changes:
  Require the new PACKET_KEEPALIVE_OFFLOAD permission.
  Add an error code for generic hardware error.
  Fix bugs and crashes in PacketKeepalive API.
  Add tests for the PacketKeepalive API.
  Add a PACKET_KEEPALIVE_OFFLOAD permission.
  Use a CountDownLatch instead of sleep() in NetworkFactory tests.
  Get rid of shortSleep() in ConnectivityServiceTest.
  Make ConnectivityServiceTest a bit more readable.
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 4055836..9a2a241 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -1243,6 +1243,8 @@
 
         /** The hardware does not support this request. */
         public static final int ERROR_HARDWARE_UNSUPPORTED = -30;
+        /** The hardware returned an error. */
+        public static final int ERROR_HARDWARE_ERROR = -31;
 
         public static final int NATT_PORT = 4500;
 
diff --git a/core/java/android/net/NetworkFactory.java b/core/java/android/net/NetworkFactory.java
index 5f46c73..cab88b9 100644
--- a/core/java/android/net/NetworkFactory.java
+++ b/core/java/android/net/NetworkFactory.java
@@ -169,7 +169,8 @@
         }
     }
 
-    private void handleAddRequest(NetworkRequest request, int score) {
+    @VisibleForTesting
+    protected void handleAddRequest(NetworkRequest request, int score) {
         NetworkRequestInfo n = mNetworkRequests.get(request.requestId);
         if (n == null) {
             if (DBG) log("got request " + request + " with score " + score);
@@ -184,7 +185,8 @@
         evalRequest(n);
     }
 
-    private void handleRemoveRequest(NetworkRequest request) {
+    @VisibleForTesting
+    protected void handleRemoveRequest(NetworkRequest request) {
         NetworkRequestInfo n = mNetworkRequests.get(request.requestId);
         if (n != null) {
             mNetworkRequests.remove(request.requestId);
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index c8bb675..b8b6444 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -1056,6 +1056,11 @@
     <permission android:name="android.permission.CONNECTIVITY_INTERNAL"
         android:protectionLevel="signature|privileged" />
 
+    <!-- Allows a system application to access hardware packet offload capabilities.
+         @hide -->
+    <permission android:name="android.permission.PACKET_KEEPALIVE_OFFLOAD"
+        android:protectionLevel="signature|privileged" />
+
     <!-- @SystemApi
          @hide -->
     <permission android:name="android.permission.RECEIVE_DATA_ACTIVITY_CHANGE"
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 1f4427f..d1e1683 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -359,6 +359,9 @@
      */
     private static final int EVENT_REGISTER_NETWORK_LISTENER_WITH_INTENT = 31;
 
+    /** Handler thread used for both of the handlers below. */
+    @VisibleForTesting
+    protected final HandlerThread mHandlerThread;
     /** Handler used for internal events. */
     final private InternalHandler mHandler;
     /** Handler used for incoming {@link NetworkStateTracker} events. */
@@ -614,6 +617,11 @@
     }
     private LegacyTypeTracker mLegacyTypeTracker = new LegacyTypeTracker();
 
+    @VisibleForTesting
+    protected HandlerThread createHandlerThread() {
+        return new HandlerThread("ConnectivityServiceThread");
+    }
+
     public ConnectivityService(Context context, INetworkManagementService netManager,
             INetworkStatsService statsService, INetworkPolicyManager policyManager) {
         if (DBG) log("ConnectivityService starting up");
@@ -627,10 +635,10 @@
         mDefaultMobileDataRequest = createInternetRequestForTransport(
                 NetworkCapabilities.TRANSPORT_CELLULAR);
 
-        HandlerThread handlerThread = new HandlerThread("ConnectivityServiceThread");
-        handlerThread.start();
-        mHandler = new InternalHandler(handlerThread.getLooper());
-        mTrackerHandler = new NetworkStateTrackerHandler(handlerThread.getLooper());
+        mHandlerThread = createHandlerThread();
+        mHandlerThread.start();
+        mHandler = new InternalHandler(mHandlerThread.getLooper());
+        mTrackerHandler = new NetworkStateTrackerHandler(mHandlerThread.getLooper());
 
         // setup our unique device name
         if (TextUtils.isEmpty(SystemProperties.get("net.hostname"))) {
@@ -1458,7 +1466,7 @@
     }
 
     private void enforceKeepalivePermission() {
-        mContext.enforceCallingPermission(KeepaliveTracker.PERMISSION, "ConnectivityService");
+        mContext.enforceCallingOrSelfPermission(KeepaliveTracker.PERMISSION, "ConnectivityService");
     }
 
     public void sendConnectedBroadcast(NetworkInfo info) {
diff --git a/services/core/java/com/android/server/connectivity/KeepalivePacketData.java b/services/core/java/com/android/server/connectivity/KeepalivePacketData.java
index 64b9399..2ccfdd1 100644
--- a/services/core/java/com/android/server/connectivity/KeepalivePacketData.java
+++ b/services/core/java/com/android/server/connectivity/KeepalivePacketData.java
@@ -71,6 +71,7 @@
         // Check we have two IP addresses of the same family.
         if (srcAddress == null || dstAddress == null ||
                 !srcAddress.getClass().getName().equals(dstAddress.getClass().getName())) {
+            throw new InvalidPacketException(ERROR_INVALID_IP_ADDRESS);
         }
 
         // Set the protocol.
@@ -102,7 +103,7 @@
             InetAddress srcAddress, int srcPort,
             InetAddress dstAddress, int dstPort) throws InvalidPacketException {
 
-        if (!(srcAddress instanceof Inet4Address)) {
+        if (!(srcAddress instanceof Inet4Address) || !(dstAddress instanceof Inet4Address)) {
             throw new InvalidPacketException(ERROR_INVALID_IP_ADDRESS);
         }
 
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index c78f347..90c9ddf 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -62,8 +62,7 @@
     private static final String TAG = "KeepaliveTracker";
     private static final boolean DBG = true;
 
-    // TODO: Change this to a system-only permission.
-    public static final String PERMISSION = android.Manifest.permission.CHANGE_NETWORK_STATE;
+    public static final String PERMISSION = android.Manifest.permission.PACKET_KEEPALIVE_OFFLOAD;
 
     /** Keeps track of keepalive requests. */
     private final HashMap <NetworkAgentInfo, HashMap<Integer, KeepaliveInfo>> mKeepalives =
@@ -208,6 +207,8 @@
                 Log.d(TAG, "Stopping keepalive " + mSlot + " on " + mNai.name());
                 mNai.asyncChannel.sendMessage(CMD_STOP_PACKET_KEEPALIVE, mSlot);
             }
+            // TODO: at the moment we unconditionally return failure here. In cases where the
+            // NetworkAgent is alive, should we ask it to reply, so it can return failure?
             notifyMessenger(mSlot, reason);
             unlinkDeathRecipient();
         }
@@ -233,17 +234,14 @@
             mKeepalives.put(nai, networkKeepalives);
         }
 
-        // Find the lowest-numbered free slot.
+        // Find the lowest-numbered free slot. Slot numbers start from 1, because that's what two
+        // separate chipset implementations independently came up with.
         int slot;
-        for (slot = 0; slot < networkKeepalives.size(); slot++) {
+        for (slot = 1; slot <= networkKeepalives.size(); slot++) {
             if (networkKeepalives.get(slot) == null) {
                 return slot;
             }
         }
-        // No free slot, pick one at the end.
-
-        // HACK for broadcom hardware that does not support slot 0!
-        if (slot == 0) slot = 1;
         return slot;
     }
 
@@ -267,14 +265,15 @@
     }
 
     public void handleStopKeepalive(NetworkAgentInfo nai, int slot, int reason) {
+        String networkName = (nai == null) ? "(null)" : nai.name();
         HashMap <Integer, KeepaliveInfo> networkKeepalives = mKeepalives.get(nai);
         if (networkKeepalives == null) {
-            Log.e(TAG, "Attempt to stop keepalive on nonexistent network " + nai.name());
+            Log.e(TAG, "Attempt to stop keepalive on nonexistent network " + networkName);
             return;
         }
         KeepaliveInfo ki = networkKeepalives.get(slot);
         if (ki == null) {
-            Log.e(TAG, "Attempt to stop nonexistent keepalive " + slot + " on " + nai.name());
+            Log.e(TAG, "Attempt to stop nonexistent keepalive " + slot + " on " + networkName);
             return;
         }
         ki.stop(reason);
@@ -332,6 +331,11 @@
 
     public void startNattKeepalive(NetworkAgentInfo nai, int intervalSeconds, Messenger messenger,
             IBinder binder, String srcAddrString, int srcPort, String dstAddrString, int dstPort) {
+        if (nai == null) {
+            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_NETWORK);
+            return;
+        }
+
         InetAddress srcAddress, dstAddress;
         try {
             srcAddress = NetworkUtils.numericToInetAddress(srcAddrString);
diff --git a/services/tests/servicestests/AndroidManifest.xml b/services/tests/servicestests/AndroidManifest.xml
index 919293a..8bb20a6 100644
--- a/services/tests/servicestests/AndroidManifest.xml
+++ b/services/tests/servicestests/AndroidManifest.xml
@@ -39,6 +39,7 @@
     <uses-permission android:name="android.permission.MODIFY_PHONE_STATE" />
     <uses-permission android:name="android.permission.INTERNET" />
     <uses-permission android:name="android.permission.CHANGE_NETWORK_STATE" />
+    <uses-permission android:name="android.permission.PACKET_KEEPALIVE_OFFLOAD" />
 
     <application>
         <uses-library android:name="android.test.runner" />
diff --git a/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java b/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
index b4c76b7..97e16da 100644
--- a/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/ConnectivityServiceTest.java
@@ -20,35 +20,9 @@
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.ConnectivityManager.TYPE_WIFI;
 import static android.net.ConnectivityManager.getNetworkTypeName;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_CBS;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_DUN;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_EIMS;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_FOTA;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_IA;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_IMS;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_MMS;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_RCS;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_SUPL;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_WIFI_P2P;
-import static android.net.NetworkCapabilities.NET_CAPABILITY_XCAP;
-import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
-import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Matchers.isA;
-import static org.mockito.Mockito.doNothing;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.doThrow;
+import static android.net.NetworkCapabilities.*;
+
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.verify;
 
 import android.app.PendingIntent;
 import android.content.BroadcastReceiver;
@@ -58,8 +32,12 @@
 import android.content.IntentFilter;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
+import android.net.ConnectivityManager.PacketKeepalive;
+import android.net.ConnectivityManager.PacketKeepaliveCallback;
 import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
+import android.net.IpPrefix;
+import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkAgent;
@@ -74,8 +52,11 @@
 import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.HandlerThread;
-import android.os.Looper;
 import android.os.INetworkManagementService;
+import android.os.Looper;
+import android.os.Message;
+import android.os.MessageQueue;
+import android.os.MessageQueue.IdleHandler;
 import android.test.AndroidTestCase;
 import android.test.suitebuilder.annotation.LargeTest;
 import android.util.Log;
@@ -84,10 +65,10 @@
 import com.android.server.connectivity.NetworkAgentInfo;
 import com.android.server.connectivity.NetworkMonitor;
 
-import org.mockito.ArgumentCaptor;
-
 import java.net.InetAddress;
-import java.util.concurrent.Future;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -99,24 +80,7 @@
 public class ConnectivityServiceTest extends AndroidTestCase {
     private static final String TAG = "ConnectivityServiceTest";
 
-    private static final String MOBILE_IFACE = "rmnet3";
-    private static final String WIFI_IFACE = "wlan6";
-
-    private static final RouteInfo MOBILE_ROUTE_V4 = RouteInfo.makeHostRoute(parse("10.0.0.33"),
-                                                                             MOBILE_IFACE);
-    private static final RouteInfo MOBILE_ROUTE_V6 = RouteInfo.makeHostRoute(parse("fd00::33"),
-                                                                             MOBILE_IFACE);
-
-    private static final RouteInfo WIFI_ROUTE_V4 = RouteInfo.makeHostRoute(parse("192.168.0.66"),
-                                                                           parse("192.168.0.1"),
-                                                                           WIFI_IFACE);
-    private static final RouteInfo WIFI_ROUTE_V6 = RouteInfo.makeHostRoute(parse("fd00::66"),
-                                                                           parse("fd00::"),
-                                                                           WIFI_IFACE);
-
-    private INetworkManagementService mNetManager;
-    private INetworkStatsService mStatsService;
-    private INetworkPolicyManager mPolicyService;
+    private static final int TIMEOUT_MS = 500;
 
     private BroadcastInterceptingContext mServiceContext;
     private WrappedConnectivityService mService;
@@ -148,14 +112,93 @@
         }
     }
 
+    /**
+     * A subclass of HandlerThread that allows callers to wait for it to become idle. waitForIdle
+     * will return immediately if the handler is already idle.
+     */
+    private class IdleableHandlerThread extends HandlerThread {
+        private IdleHandler mIdleHandler;
+
+        public IdleableHandlerThread(String name) {
+            super(name);
+        }
+
+        public void waitForIdle(int timeoutMs) {
+            final ConditionVariable cv = new ConditionVariable();
+            final MessageQueue queue = getLooper().getQueue();
+
+            synchronized (queue) {
+                if (queue.isIdle()) {
+                    return;
+                }
+
+                assertNull("BUG: only one idle handler allowed", mIdleHandler);
+                mIdleHandler = new IdleHandler() {
+                    public boolean queueIdle() {
+                        cv.open();
+                        mIdleHandler = null;
+                        return false;  // Remove the handler.
+                    }
+                };
+                queue.addIdleHandler(mIdleHandler);
+            }
+
+            if (!cv.block(timeoutMs)) {
+                fail("HandlerThread " + getName() +
+                        " did not become idle after " + timeoutMs + " ms");
+                queue.removeIdleHandler(mIdleHandler);
+            }
+        }
+    }
+
+    // Tests that IdleableHandlerThread works as expected.
+    public void testIdleableHandlerThread() {
+        final int attempts = 50;  // Causes the test to take about 200ms on bullhead-eng.
+
+        // Tests that waitForIdle returns immediately if the service is already idle.
+        for (int i = 0; i < attempts; i++) {
+            mService.waitForIdle();
+        }
+
+        // Bring up a network that we can use to send messages to ConnectivityService.
+        ConditionVariable cv = waitForConnectivityBroadcasts(1);
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.connect(false);
+        waitFor(cv);
+        Network n = mWiFiNetworkAgent.getNetwork();
+        assertNotNull(n);
+
+        // Tests that calling waitForIdle waits for messages to be processed.
+        for (int i = 0; i < attempts; i++) {
+            mWiFiNetworkAgent.setSignalStrength(i);
+            mService.waitForIdle();
+            assertEquals(i, mCm.getNetworkCapabilities(n).getSignalStrength());
+        }
+
+        // Ensure that not calling waitForIdle causes a race condition.
+        for (int i = 0; i < attempts; i++) {
+            mWiFiNetworkAgent.setSignalStrength(i);
+            if (i != mCm.getNetworkCapabilities(n).getSignalStrength()) {
+                // We hit a race condition, as expected. Pass the test.
+                return;
+            }
+        }
+
+        // No race? There is a bug in this test.
+        fail("expected race condition at least once in " + attempts + " attempts");
+    }
+
     private class MockNetworkAgent {
         private final WrappedNetworkMonitor mWrappedNetworkMonitor;
         private final NetworkInfo mNetworkInfo;
         private final NetworkCapabilities mNetworkCapabilities;
-        private final Thread mThread;
+        private final IdleableHandlerThread mHandlerThread;
         private final ConditionVariable mDisconnected = new ConditionVariable();
         private int mScore;
         private NetworkAgent mNetworkAgent;
+        private int mStartKeepaliveError = PacketKeepalive.ERROR_HARDWARE_UNSUPPORTED;
+        private int mStopKeepaliveError = PacketKeepalive.NO_KEEPALIVE;
+        private Integer mExpectedKeepaliveSlot = null;
 
         MockNetworkAgent(int transport) {
             final int type = transportToLegacyType(transport);
@@ -173,26 +216,42 @@
                 default:
                     throw new UnsupportedOperationException("unimplemented network type");
             }
-            final ConditionVariable initComplete = new ConditionVariable();
-            final ConditionVariable networkMonitorAvailable = mService.getNetworkMonitorCreatedCV();
-            mThread = new Thread() {
-                public void run() {
-                    Looper.prepare();
-                    mNetworkAgent = new NetworkAgent(Looper.myLooper(), mServiceContext,
-                            "Mock" + typeName, mNetworkInfo, mNetworkCapabilities,
-                            new LinkProperties(), mScore, new NetworkMisc()) {
-                        public void unwanted() { mDisconnected.open(); }
-                    };
-                    initComplete.open();
-                    Looper.loop();
+            mHandlerThread = new IdleableHandlerThread("Mock-" + typeName);
+            mHandlerThread.start();
+            mNetworkAgent = new NetworkAgent(mHandlerThread.getLooper(), mServiceContext,
+                    "Mock-" + typeName, mNetworkInfo, mNetworkCapabilities,
+                    new LinkProperties(), mScore, new NetworkMisc()) {
+                @Override
+                public void unwanted() { mDisconnected.open(); }
+
+                @Override
+                public void startPacketKeepalive(Message msg) {
+                    int slot = msg.arg1;
+                    if (mExpectedKeepaliveSlot != null) {
+                        assertEquals((int) mExpectedKeepaliveSlot, slot);
+                    }
+                    onPacketKeepaliveEvent(slot, mStartKeepaliveError);
+                }
+
+                @Override
+                public void stopPacketKeepalive(Message msg) {
+                    onPacketKeepaliveEvent(msg.arg1, mStopKeepaliveError);
                 }
             };
-            mThread.start();
-            waitFor(initComplete);
-            waitFor(networkMonitorAvailable);
+            // Waits for the NetworkAgent to be registered, which includes the creation of the
+            // NetworkMonitor.
+            mService.waitForIdle();
             mWrappedNetworkMonitor = mService.getLastCreatedWrappedNetworkMonitor();
         }
 
+        public void waitForIdle(int timeoutMs) {
+            mHandlerThread.waitForIdle(timeoutMs);
+        }
+
+        public void waitForIdle() {
+            waitForIdle(TIMEOUT_MS);
+        }
+
         public void adjustScore(int change) {
             mScore += change;
             mNetworkAgent.sendNetworkScore(mScore);
@@ -203,6 +262,11 @@
             mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
         }
 
+        public void setSignalStrength(int signalStrength) {
+            mNetworkCapabilities.setSignalStrength(signalStrength);
+            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+        }
+
         public void connectWithoutInternet() {
             mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
             mNetworkAgent.sendNetworkInfo(mNetworkInfo);
@@ -272,15 +336,46 @@
         public WrappedNetworkMonitor getWrappedNetworkMonitor() {
             return mWrappedNetworkMonitor;
         }
+
+        public void sendLinkProperties(LinkProperties lp) {
+            mNetworkAgent.sendLinkProperties(lp);
+        }
+
+        public void setStartKeepaliveError(int error) {
+            mStartKeepaliveError = error;
+        }
+
+        public void setStopKeepaliveError(int error) {
+            mStopKeepaliveError = error;
+        }
+
+        public void setExpectedKeepaliveSlot(Integer slot) {
+            mExpectedKeepaliveSlot = slot;
+        }
     }
 
+    /**
+     * A NetworkFactory that allows tests to wait until any in-flight NetworkRequest add or remove
+     * operations have been processed. Before ConnectivityService can add or remove any requests,
+     * the factory must be told to expect those operations by calling expectAddRequests or
+     * expectRemoveRequests.
+     */
     private static class MockNetworkFactory extends NetworkFactory {
         private final ConditionVariable mNetworkStartedCV = new ConditionVariable();
         private final ConditionVariable mNetworkStoppedCV = new ConditionVariable();
-        private final ConditionVariable mNetworkRequestedCV = new ConditionVariable();
-        private final ConditionVariable mNetworkReleasedCV = new ConditionVariable();
         private final AtomicBoolean mNetworkStarted = new AtomicBoolean(false);
 
+        // Used to expect that requests be removed or added on a separate thread, without sleeping.
+        // Callers can call either expectAddRequests() or expectRemoveRequests() exactly once, then
+        // cause some other thread to add or remove requests, then call waitForRequests(). We can
+        // either expect requests to be added or removed, but not both, because CountDownLatch can
+        // only count in one direction.
+        private CountDownLatch mExpectations;
+
+        // Whether we are currently expecting requests to be added or removed. Valid only if
+        // mExpectations is non-null.
+        private boolean mExpectingAdditions;
+
         public MockNetworkFactory(Looper looper, Context context, String logTag,
                 NetworkCapabilities filter) {
             super(looper, context, logTag, filter);
@@ -314,28 +409,75 @@
             return mNetworkStoppedCV;
         }
 
-        protected void needNetworkFor(NetworkRequest networkRequest, int score) {
-            super.needNetworkFor(networkRequest, score);
-            mNetworkRequestedCV.open();
+        @Override
+        protected void handleAddRequest(NetworkRequest request, int score) {
+            // If we're expecting anything, we must be expecting additions.
+            if (mExpectations != null && !mExpectingAdditions) {
+                fail("Can't add requests while expecting requests to be removed");
+            }
+
+            // Add the request.
+            super.handleAddRequest(request, score);
+
+            // Reduce the number of request additions we're waiting for.
+            if (mExpectingAdditions) {
+                assertTrue("Added more requests than expected", mExpectations.getCount() > 0);
+                mExpectations.countDown();
+            }
         }
 
-        protected void releaseNetworkFor(NetworkRequest networkRequest) {
-            super.releaseNetworkFor(networkRequest);
-            mNetworkReleasedCV.open();
+        @Override
+        protected void handleRemoveRequest(NetworkRequest request) {
+            // If we're expecting anything, we must be expecting removals.
+            if (mExpectations != null && mExpectingAdditions) {
+                fail("Can't remove requests while expecting requests to be added");
+            }
+
+            // Remove the request.
+            super.handleRemoveRequest(request);
+
+            // Reduce the number of request removals we're waiting for.
+            if (!mExpectingAdditions) {
+                assertTrue("Removed more requests than expected", mExpectations.getCount() > 0);
+                mExpectations.countDown();
+            }
         }
 
-        public ConditionVariable getNetworkRequestedCV() {
-            mNetworkRequestedCV.close();
-            return mNetworkRequestedCV;
+        private void assertNoExpectations() {
+            if (mExpectations != null) {
+                fail("Can't add expectation, " + mExpectations.getCount() + " already pending");
+            }
         }
 
-        public ConditionVariable getNetworkReleasedCV() {
-            mNetworkReleasedCV.close();
-            return mNetworkReleasedCV;
+        // Expects that count requests will be added.
+        public void expectAddRequests(final int count) {
+            assertNoExpectations();
+            mExpectingAdditions = true;
+            mExpectations = new CountDownLatch(count);
         }
 
-        public void waitForNetworkRequests(final int count) {
-            waitFor(new Criteria() { public boolean get() { return count == getRequestCount(); } });
+        // Expects that count requests will be removed.
+        public void expectRemoveRequests(final int count) {
+            assertNoExpectations();
+            mExpectingAdditions = false;
+            mExpectations = new CountDownLatch(count);
+        }
+
+        // Waits for the expected request additions or removals to happen within a timeout.
+        public void waitForRequests() throws InterruptedException {
+            assertNotNull("Nothing to wait for", mExpectations);
+            mExpectations.await(TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            final long count = mExpectations.getCount();
+            final String msg = count + " requests still not " +
+                    (mExpectingAdditions ? "added" : "removed") +
+                    " after " + TIMEOUT_MS + " ms";
+            assertEquals(msg, 0, count);
+            mExpectations = null;
+        }
+
+        public void waitForNetworkRequests(final int count) throws InterruptedException {
+            waitForRequests();
+            assertEquals(count, getMyRequestCount());
         }
     }
 
@@ -356,7 +498,6 @@
     }
 
     private class WrappedConnectivityService extends ConnectivityService {
-        private final ConditionVariable mNetworkMonitorCreated = new ConditionVariable();
         private WrappedNetworkMonitor mLastCreatedNetworkMonitor;
 
         public WrappedConnectivityService(Context context, INetworkManagementService netManager,
@@ -365,6 +506,11 @@
         }
 
         @Override
+        protected HandlerThread createHandlerThread() {
+            return new IdleableHandlerThread("WrappedConnectivityService");
+        }
+
+        @Override
         protected int getDefaultTcpRwnd() {
             // Prevent wrapped ConnectivityService from trying to write to SystemProperties.
             return 0;
@@ -397,7 +543,6 @@
             final WrappedNetworkMonitor monitor = new WrappedNetworkMonitor(context, handler, nai,
                     defaultRequest);
             mLastCreatedNetworkMonitor = monitor;
-            mNetworkMonitorCreated.open();
             return monitor;
         }
 
@@ -405,10 +550,14 @@
             return mLastCreatedNetworkMonitor;
         }
 
-        public ConditionVariable getNetworkMonitorCreatedCV() {
-            mNetworkMonitorCreated.close();
-            return mNetworkMonitorCreated;
+        public void waitForIdle(int timeoutMs) {
+            ((IdleableHandlerThread) mHandlerThread).waitForIdle(timeoutMs);
         }
+
+        public void waitForIdle() {
+            waitForIdle(TIMEOUT_MS);
+        }
+
     }
 
     private interface Criteria {
@@ -431,23 +580,11 @@
     }
 
     /**
-     * Wait up to 500ms for {@code conditonVariable} to open.
-     * Fails if 500ms goes by before {@code conditionVariable} opens.
+     * Wait up to TIMEOUT_MS for {@code conditionVariable} to open.
+     * Fails if TIMEOUT_MS goes by before {@code conditionVariable} opens.
      */
     static private void waitFor(ConditionVariable conditionVariable) {
-        assertTrue(conditionVariable.block(500));
-    }
-
-    /**
-     * This should only be used to verify that nothing happens, in other words that no unexpected
-     * changes occur.  It should never be used to wait for a specific positive signal to occur.
-     */
-    private void shortSleep() {
-        // TODO: Instead of sleeping, instead wait for all message loops to idle.
-        try {
-            Thread.sleep(500);
-        } catch (InterruptedException e) {
-        }
+        assertTrue(conditionVariable.block(TIMEOUT_MS));
     }
 
     @Override
@@ -455,13 +592,11 @@
         super.setUp();
 
         mServiceContext = new MockContext(getContext());
+        mService = new WrappedConnectivityService(mServiceContext,
+                mock(INetworkManagementService.class),
+                mock(INetworkStatsService.class),
+                mock(INetworkPolicyManager.class));
 
-        mNetManager = mock(INetworkManagementService.class);
-        mStatsService = mock(INetworkStatsService.class);
-        mPolicyService = mock(INetworkPolicyManager.class);
-
-        mService = new WrappedConnectivityService(
-                mServiceContext, mNetManager, mStatsService, mPolicyService);
         mService.systemReady();
         mCm = new ConnectivityManager(getContext(), mService);
     }
@@ -583,11 +718,11 @@
         // Test bringing up unvalidated cellular
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
-        shortSleep();
+        mService.waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test cellular disconnect.
         mCellNetworkAgent.disconnect();
-        shortSleep();
+        mService.waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up validated cellular
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -797,6 +932,11 @@
         LOST
     }
 
+    /**
+     * 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
+     * received. assertNoCallback may be called at any time.
+     */
     private class TestNetworkCallback extends NetworkCallback {
         private final ConditionVariable mConditionVariable = new ConditionVariable();
         private CallbackState mLastCallback = CallbackState.NONE;
@@ -819,14 +959,15 @@
             mConditionVariable.open();
         }
 
-        ConditionVariable getConditionVariable() {
+        void expectCallback(CallbackState state) {
+            waitFor(mConditionVariable);
+            assertEquals(state, mLastCallback);
             mLastCallback = CallbackState.NONE;
             mConditionVariable.close();
-            return mConditionVariable;
         }
 
-        CallbackState getLastCallback() {
-            return mLastCallback;
+        void assertNoCallback() {
+            assertEquals(CallbackState.NONE, mLastCallback);
         }
     }
 
@@ -842,98 +983,68 @@
         mCm.registerNetworkCallback(cellRequest, cellNetworkCallback);
 
         // Test unvalidated networks
-        ConditionVariable cellCv = cellNetworkCallback.getConditionVariable();
-        ConditionVariable wifiCv = wifiNetworkCallback.getConditionVariable();
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
-        waitFor(cellCv);
-        assertEquals(CallbackState.AVAILABLE, cellNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
+        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE);
+        wifiNetworkCallback.assertNoCallback();
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         waitFor(cv);
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         // This should not trigger spurious onAvailable() callbacks, b/21762680.
         mCellNetworkAgent.adjustScore(-1);
-        shortSleep();
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, cellNetworkCallback.getLastCallback());
+        mService.waitForIdle();
+        wifiNetworkCallback.assertNoCallback();
+        cellNetworkCallback.assertNoCallback();
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
-        waitFor(wifiCv);
-        assertEquals(CallbackState.AVAILABLE, wifiNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, cellNetworkCallback.getLastCallback());
+        wifiNetworkCallback.expectCallback(CallbackState.AVAILABLE);
+        cellNetworkCallback.assertNoCallback();
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         waitFor(cv);
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.disconnect();
-        waitFor(wifiCv);
-        assertEquals(CallbackState.LOST, wifiNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, cellNetworkCallback.getLastCallback());
+        wifiNetworkCallback.expectCallback(CallbackState.LOST);
+        cellNetworkCallback.assertNoCallback();
         waitFor(cv);
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.disconnect();
-        waitFor(cellCv);
-        assertEquals(CallbackState.LOST, cellNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
+        cellNetworkCallback.expectCallback(CallbackState.LOST);
+        wifiNetworkCallback.assertNoCallback();
         waitFor(cv);
 
         // Test validated networks
-
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
-        waitFor(cellCv);
-        assertEquals(CallbackState.AVAILABLE, cellNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
+        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE);
+        wifiNetworkCallback.assertNoCallback();
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         // This should not trigger spurious onAvailable() callbacks, b/21762680.
         mCellNetworkAgent.adjustScore(-1);
-        shortSleep();
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, cellNetworkCallback.getLastCallback());
+        mService.waitForIdle();
+        wifiNetworkCallback.assertNoCallback();
+        cellNetworkCallback.assertNoCallback();
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
-        waitFor(wifiCv);
-        assertEquals(CallbackState.AVAILABLE, wifiNetworkCallback.getLastCallback());
-        waitFor(cellCv);
-        assertEquals(CallbackState.LOSING, cellNetworkCallback.getLastCallback());
+        wifiNetworkCallback.expectCallback(CallbackState.AVAILABLE);
+        cellNetworkCallback.expectCallback(CallbackState.LOSING);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         mWiFiNetworkAgent.disconnect();
-        waitFor(wifiCv);
-        assertEquals(CallbackState.LOST, wifiNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, cellNetworkCallback.getLastCallback());
+        wifiNetworkCallback.expectCallback(CallbackState.LOST);
+        cellNetworkCallback.assertNoCallback();
 
-        cellCv = cellNetworkCallback.getConditionVariable();
-        wifiCv = wifiNetworkCallback.getConditionVariable();
         mCellNetworkAgent.disconnect();
-        waitFor(cellCv);
-        assertEquals(CallbackState.LOST, cellNetworkCallback.getLastCallback());
-        assertEquals(CallbackState.NONE, wifiNetworkCallback.getLastCallback());
+        cellNetworkCallback.expectCallback(CallbackState.LOST);
+        wifiNetworkCallback.assertNoCallback();
     }
 
     private void tryNetworkFactoryRequests(int capability) throws Exception {
@@ -957,18 +1068,21 @@
                 mServiceContext, "testFactory", filter);
         testFactory.setScoreFilter(40);
         ConditionVariable cv = testFactory.getNetworkStartedCV();
+        testFactory.expectAddRequests(1);
         testFactory.register();
+        testFactory.waitForNetworkRequests(1);
         int expectedRequestCount = 1;
         NetworkCallback networkCallback = null;
         // For non-INTERNET capabilities we cannot rely on the default request being present, so
         // add one.
         if (capability != NET_CAPABILITY_INTERNET) {
-            testFactory.waitForNetworkRequests(1);
             assertFalse(testFactory.getMyStartRequested());
             NetworkRequest request = new NetworkRequest.Builder().addCapability(capability).build();
             networkCallback = new NetworkCallback();
+            testFactory.expectAddRequests(1);
             mCm.requestNetwork(request, networkCallback);
             expectedRequestCount++;
+            testFactory.waitForNetworkRequests(expectedRequestCount);
         }
         waitFor(cv);
         assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
@@ -981,13 +1095,20 @@
         // unvalidated penalty.
         testAgent.adjustScore(40);
         cv = testFactory.getNetworkStoppedCV();
+
+        // When testAgent connects, ConnectivityService will re-send us all current requests with
+        // the new score. There are expectedRequestCount such requests, and we must wait for all of
+        // them.
+        testFactory.expectAddRequests(expectedRequestCount);
         testAgent.connect(false);
         testAgent.addCapability(capability);
         waitFor(cv);
-        assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
+        testFactory.waitForNetworkRequests(expectedRequestCount);
         assertFalse(testFactory.getMyStartRequested());
 
         // Bring in a bunch of requests.
+        testFactory.expectAddRequests(10);
+        assertEquals(expectedRequestCount, testFactory.getMyRequestCount());
         ConnectivityManager.NetworkCallback[] networkCallbacks =
                 new ConnectivityManager.NetworkCallback[10];
         for (int i = 0; i< networkCallbacks.length; i++) {
@@ -1000,6 +1121,7 @@
         assertFalse(testFactory.getMyStartRequested());
 
         // Remove the requests.
+        testFactory.expectRemoveRequests(10);
         for (int i = 0; i < networkCallbacks.length; i++) {
             mCm.unregisterNetworkCallback(networkCallbacks[i]);
         }
@@ -1088,10 +1210,8 @@
         // Test bringing up unvalidated cellular with MMS
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.addCapability(NET_CAPABILITY_MMS);
-        cv = networkCallback.getConditionVariable();
         mCellNetworkAgent.connectWithoutInternet();
-        waitFor(cv);
-        assertEquals(CallbackState.AVAILABLE, networkCallback.getLastCallback());
+        networkCallback.expectCallback(CallbackState.AVAILABLE);
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test releasing NetworkRequest disconnects cellular with MMS
         cv = mCellNetworkAgent.getDisconnectedCV();
@@ -1114,12 +1234,10 @@
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         mCm.requestNetwork(builder.build(), networkCallback);
         // Test bringing up MMS cellular network
-        cv = networkCallback.getConditionVariable();
         MockNetworkAgent mmsNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mmsNetworkAgent.addCapability(NET_CAPABILITY_MMS);
         mmsNetworkAgent.connectWithoutInternet();
-        waitFor(cv);
-        assertEquals(CallbackState.AVAILABLE, networkCallback.getLastCallback());
+        networkCallback.expectCallback(CallbackState.AVAILABLE);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
         // Test releasing MMS NetworkRequest does not disconnect main cellular NetworkAgent
         cv = mmsNetworkAgent.getDisconnectedCV();
@@ -1139,133 +1257,245 @@
         final NetworkRequest validatedRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_VALIDATED).build();
         mCm.registerNetworkCallback(validatedRequest, validatedCallback);
-        ConditionVariable validatedCv = validatedCallback.getConditionVariable();
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
-        ConditionVariable cv = captivePortalCallback.getConditionVariable();
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithCaptivePortal();
-        waitFor(cv);
-        assertEquals(CallbackState.AVAILABLE, captivePortalCallback.getLastCallback());
+        captivePortalCallback.expectCallback(CallbackState.AVAILABLE);
 
         // Take down network.
         // Expect onLost callback.
-        cv = captivePortalCallback.getConditionVariable();
         mWiFiNetworkAgent.disconnect();
-        waitFor(cv);
-        assertEquals(CallbackState.LOST, captivePortalCallback.getLastCallback());
+        captivePortalCallback.expectCallback(CallbackState.LOST);
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
-        cv = captivePortalCallback.getConditionVariable();
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithCaptivePortal();
-        waitFor(cv);
-        assertEquals(CallbackState.AVAILABLE, captivePortalCallback.getLastCallback());
+        captivePortalCallback.expectCallback(CallbackState.AVAILABLE);
 
         // Make captive portal disappear then revalidate.
         // Expect onLost callback because network no longer provides NET_CAPABILITY_CAPTIVE_PORTAL.
-        cv = captivePortalCallback.getConditionVariable();
         mWiFiNetworkAgent.getWrappedNetworkMonitor().gen204ProbeResult = 204;
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        waitFor(cv);
-        assertEquals(CallbackState.LOST, captivePortalCallback.getLastCallback());
+        captivePortalCallback.expectCallback(CallbackState.LOST);
 
         // Expect NET_CAPABILITY_VALIDATED onAvailable callback.
-        waitFor(validatedCv);
-        assertEquals(CallbackState.AVAILABLE, validatedCallback.getLastCallback());
+        validatedCallback.expectCallback(CallbackState.AVAILABLE);
 
         // Break network connectivity.
         // Expect NET_CAPABILITY_VALIDATED onLost callback.
-        validatedCv = validatedCallback.getConditionVariable();
         mWiFiNetworkAgent.getWrappedNetworkMonitor().gen204ProbeResult = 500;
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
-        waitFor(validatedCv);
-        assertEquals(CallbackState.LOST, validatedCallback.getLastCallback());
+        validatedCallback.expectCallback(CallbackState.LOST);
     }
 
-//    @Override
-//    public void tearDown() throws Exception {
-//        super.tearDown();
-//    }
-//
-//    public void testMobileConnectedAddedRoutes() throws Exception {
-//        Future<?> nextConnBroadcast;
-//
-//        // bring up mobile network
-//        mMobile.info.setDetailedState(DetailedState.CONNECTED, null, null);
-//        mMobile.link.setInterfaceName(MOBILE_IFACE);
-//        mMobile.link.addRoute(MOBILE_ROUTE_V4);
-//        mMobile.link.addRoute(MOBILE_ROUTE_V6);
-//        mMobile.doReturnDefaults();
-//
-//        cv = waitForConnectivityBroadcasts(1);
-//        mTrackerHandler.obtainMessage(EVENT_STATE_CHANGED, mMobile.info).sendToTarget();
-//        waitFor(cv);
-//
-//        // verify that both routes were added
-//        int mobileNetId = mMobile.tracker.getNetwork().netId;
-//        verify(mNetManager).addRoute(eq(mobileNetId), eq(MOBILE_ROUTE_V4));
-//        verify(mNetManager).addRoute(eq(mobileNetId), eq(MOBILE_ROUTE_V6));
-//    }
-//
-//    public void testMobileWifiHandoff() throws Exception {
-//        Future<?> nextConnBroadcast;
-//
-//        // bring up mobile network
-//        mMobile.info.setDetailedState(DetailedState.CONNECTED, null, null);
-//        mMobile.link.setInterfaceName(MOBILE_IFACE);
-//        mMobile.link.addRoute(MOBILE_ROUTE_V4);
-//        mMobile.link.addRoute(MOBILE_ROUTE_V6);
-//        mMobile.doReturnDefaults();
-//
-//        cv = waitForConnectivityBroadcasts(1);
-//        mTrackerHandler.obtainMessage(EVENT_STATE_CHANGED, mMobile.info).sendToTarget();
-//        waitFor(cv);
-//
-//        reset(mNetManager);
-//
-//        // now bring up wifi network
-//        mWifi.info.setDetailedState(DetailedState.CONNECTED, null, null);
-//        mWifi.link.setInterfaceName(WIFI_IFACE);
-//        mWifi.link.addRoute(WIFI_ROUTE_V4);
-//        mWifi.link.addRoute(WIFI_ROUTE_V6);
-//        mWifi.doReturnDefaults();
-//
-//        // expect that mobile will be torn down
-//        doReturn(true).when(mMobile.tracker).teardown();
-//
-//        cv = waitForConnectivityBroadcasts(1);
-//        mTrackerHandler.obtainMessage(EVENT_STATE_CHANGED, mWifi.info).sendToTarget();
-//        waitFor(cv);
-//
-//        // verify that wifi routes added, and teardown requested
-//        int wifiNetId = mWifi.tracker.getNetwork().netId;
-//        verify(mNetManager).addRoute(eq(wifiNetId), eq(WIFI_ROUTE_V4));
-//        verify(mNetManager).addRoute(eq(wifiNetId), eq(WIFI_ROUTE_V6));
-//        verify(mMobile.tracker).teardown();
-//
-//        int mobileNetId = mMobile.tracker.getNetwork().netId;
-//
-//        reset(mNetManager, mMobile.tracker);
-//
-//        // tear down mobile network, as requested
-//        mMobile.info.setDetailedState(DetailedState.DISCONNECTED, null, null);
-//        mMobile.link.clear();
-//        mMobile.doReturnDefaults();
-//
-//        cv = waitForConnectivityBroadcasts(1);
-//        mTrackerHandler.obtainMessage(EVENT_STATE_CHANGED, mMobile.info).sendToTarget();
-//        waitFor(cv);
-//
-//        verify(mNetManager).removeRoute(eq(mobileNetId), eq(MOBILE_ROUTE_V4));
-//        verify(mNetManager).removeRoute(eq(mobileNetId), eq(MOBILE_ROUTE_V6));
-//
-//    }
+    private static class TestKeepaliveCallback extends PacketKeepaliveCallback {
 
-    private static InetAddress parse(String addr) {
-        return InetAddress.parseNumericAddress(addr);
+        public static enum CallbackType { ON_STARTED, ON_STOPPED, ON_ERROR };
+
+        private class CallbackValue {
+            public CallbackType callbackType;
+            public int error;
+
+            public CallbackValue(CallbackType type) {
+                this.callbackType = type;
+                this.error = PacketKeepalive.SUCCESS;
+                assertTrue("onError callback must have error", type != CallbackType.ON_ERROR);
+            }
+
+            public CallbackValue(CallbackType type, int error) {
+                this.callbackType = type;
+                this.error = error;
+                assertEquals("error can only be set for onError", type, CallbackType.ON_ERROR);
+            }
+
+            @Override
+            public boolean equals(Object o) {
+                return o instanceof CallbackValue &&
+                        this.callbackType == ((CallbackValue) o).callbackType &&
+                        this.error == ((CallbackValue) o).error;
+            }
+
+            @Override
+            public String toString() {
+                return String.format("%s(%s, %d)", getClass().getSimpleName(), callbackType, error);
+            }
+        }
+
+        private LinkedBlockingQueue<CallbackValue> mCallbacks = new LinkedBlockingQueue<>();
+
+        @Override
+        public void onStarted() {
+            mCallbacks.add(new CallbackValue(CallbackType.ON_STARTED));
+        }
+
+        @Override
+        public void onStopped() {
+            mCallbacks.add(new CallbackValue(CallbackType.ON_STOPPED));
+        }
+
+        @Override
+        public void onError(int error) {
+            mCallbacks.add(new CallbackValue(CallbackType.ON_ERROR, error));
+        }
+
+        private void expectCallback(CallbackValue callbackValue) {
+            try {
+                assertEquals(
+                        callbackValue,
+                        mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+            } catch (InterruptedException e) {
+                fail(callbackValue.callbackType + " callback not seen after " + TIMEOUT_MS + " ms");
+            }
+        }
+
+        public void expectStarted() {
+            expectCallback(new CallbackValue(CallbackType.ON_STARTED));
+        }
+
+        public void expectStopped() {
+            expectCallback(new CallbackValue(CallbackType.ON_STOPPED));
+        }
+
+        public void expectError(int error) {
+            expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
+        }
     }
 
+    private Network connectKeepaliveNetwork(LinkProperties lp) {
+        // Ensure the network is disconnected before we do anything.
+        if (mWiFiNetworkAgent != null) {
+            assertNull(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()));
+        }
+
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        ConditionVariable cv = waitForConnectivityBroadcasts(1);
+        mWiFiNetworkAgent.connect(true);
+        waitFor(cv);
+        verifyActiveNetwork(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.sendLinkProperties(lp);
+        mService.waitForIdle();
+        return mWiFiNetworkAgent.getNetwork();
+    }
+
+    public void testPacketKeepalives() throws Exception {
+        InetAddress myIPv4 = InetAddress.getByName("192.0.2.129");
+        InetAddress notMyIPv4 = InetAddress.getByName("192.0.2.35");
+        InetAddress myIPv6 = InetAddress.getByName("2001:db8::1");
+        InetAddress dstIPv4 = InetAddress.getByName("8.8.8.8");
+        InetAddress dstIPv6 = InetAddress.getByName("2001:4860:4860::8888");
+
+        LinkProperties lp = new LinkProperties();
+        lp.setInterfaceName("wlan12");
+        lp.addLinkAddress(new LinkAddress(myIPv6, 64));
+        lp.addLinkAddress(new LinkAddress(myIPv4, 25));
+        lp.addRoute(new RouteInfo(InetAddress.getByName("fe80::1234")));
+        lp.addRoute(new RouteInfo(InetAddress.getByName("192.0.2.254")));
+
+        Network notMyNet = new Network(61234);
+        Network myNet = connectKeepaliveNetwork(lp);
+
+        TestKeepaliveCallback callback = new TestKeepaliveCallback();
+        PacketKeepalive ka;
+
+        // Attempt to start keepalives with invalid parameters and check for errors.
+        ka = mCm.startNattKeepalive(notMyNet, 25, callback, myIPv4, 1234, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_NETWORK);
+
+        ka = mCm.startNattKeepalive(myNet, 19, callback, notMyIPv4, 1234, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_INTERVAL);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 1234, dstIPv6);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_IP_ADDRESS);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv6, 1234, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_IP_ADDRESS);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv6, 1234, dstIPv6);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_IP_ADDRESS);  // NAT-T is IPv4-only.
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 123456, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_PORT);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 123456, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_PORT);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_HARDWARE_UNSUPPORTED);
+
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectError(PacketKeepalive.ERROR_HARDWARE_UNSUPPORTED);
+
+        // Check that a started keepalive can be stopped.
+        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectStarted();
+        mWiFiNetworkAgent.setStopKeepaliveError(PacketKeepalive.SUCCESS);
+        ka.stop();
+        callback.expectStopped();
+
+        // Check that deleting the IP address stops the keepalive.
+        LinkProperties bogusLp = new LinkProperties(lp);
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectStarted();
+        bogusLp.removeLinkAddress(new LinkAddress(myIPv4, 25));
+        bogusLp.addLinkAddress(new LinkAddress(notMyIPv4, 25));
+        mWiFiNetworkAgent.sendLinkProperties(bogusLp);
+        callback.expectError(PacketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        mWiFiNetworkAgent.sendLinkProperties(lp);
+
+        // Check that a started keepalive is stopped correctly when the network disconnects.
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectStarted();
+        mWiFiNetworkAgent.disconnect();
+        callback.expectError(PacketKeepalive.ERROR_INVALID_NETWORK);
+
+        // ... and that stopping it after that has no adverse effects.
+        assertNull(mCm.getNetworkCapabilities(myNet));
+        ka.stop();
+
+        // Reconnect.
+        myNet = connectKeepaliveNetwork(lp);
+        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
+
+        // Check things work as expected when the keepalive is stopped and the network disconnects.
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectStarted();
+        ka.stop();
+        mWiFiNetworkAgent.disconnect();
+        mService.waitForIdle();
+        callback.expectStopped();
+
+        // Reconnect.
+        myNet = connectKeepaliveNetwork(lp);
+        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
+
+        // Check that keepalive slots start from 1 and increment. The first one gets slot 1.
+        mWiFiNetworkAgent.setExpectedKeepaliveSlot(1);
+        ka = mCm.startNattKeepalive(myNet, 25, callback, myIPv4, 12345, dstIPv4);
+        callback.expectStarted();
+
+        // The second one gets slot 2.
+        mWiFiNetworkAgent.setExpectedKeepaliveSlot(2);
+        TestKeepaliveCallback callback2 = new TestKeepaliveCallback();
+        PacketKeepalive ka2 = mCm.startNattKeepalive(myNet, 25, callback2, myIPv4, 6789, dstIPv4);
+        callback2.expectStarted();
+
+        // Now stop the first one and create a third. This also gets slot 1.
+        ka.stop();
+        callback.expectStopped();
+
+        mWiFiNetworkAgent.setExpectedKeepaliveSlot(1);
+        TestKeepaliveCallback callback3 = new TestKeepaliveCallback();
+        PacketKeepalive ka3 = mCm.startNattKeepalive(myNet, 25, callback3, myIPv4, 9876, dstIPv4);
+        callback3.expectStarted();
+
+        ka2.stop();
+        callback2.expectStopped();
+
+        ka3.stop();
+        callback3.expectStopped();
+    }
 }