Merge changes I8fcaa1f0,I8925fbe4,I78394645,Iadf9f060

* changes:
  FinalizePausedKeepalive in handleStopAllKeepalives.
  Add tests for pause, resume and stop keepalives.
  Ensure autoKi is not stored when keepalive stops or is not started.
  Add tests for NATT keepalives stopped internally in KeepaliveTracker.
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index f8285ed..62d79a3 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -18,6 +18,7 @@
 
 import static android.net.SocketKeepalive.ERROR_INVALID_SOCKET;
 import static android.net.SocketKeepalive.MIN_INTERVAL_SEC;
+import static android.net.SocketKeepalive.SUCCESS;
 import static android.net.SocketKeepalive.SUCCESS_PAUSED;
 import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 import static android.system.OsConstants.AF_INET;
@@ -52,6 +53,7 @@
 import android.system.StructTimeval;
 import android.util.LocalLog;
 import android.util.Log;
+import android.util.Pair;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -381,7 +383,11 @@
             return;
         }
         autoKi.mAutomaticOnOffState = STATE_ENABLED;
-        handleResumeKeepalive(newKi);
+        final int error = handleResumeKeepalive(newKi);
+        if (error != SUCCESS) {
+            // Failed to start the keepalive
+            cleanupAutoOnOffKeepalive(autoKi);
+        }
     }
 
     /**
@@ -402,7 +408,20 @@
      * Forward to KeepaliveTracker.
      */
     public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
-        mKeepaliveTracker.handleEventSocketKeepalive(nai, slot, reason);
+        if (mKeepaliveTracker.handleEventSocketKeepalive(nai, slot, reason)) return;
+
+        // The keepalive was stopped and so the autoKi should be cleaned up.
+        final AutomaticOnOffKeepalive autoKi =
+                CollectionUtils.findFirst(
+                        mAutomaticOnOffKeepalives, it -> it.match(nai.network(), slot));
+        if (autoKi == null) {
+            // This may occur when the autoKi gets cleaned up elsewhere (i.e
+            // handleCheckKeepalivesStillValid) while waiting for the network agent to
+            // start the keepalive and the network agent returns an error event.
+            Log.e(TAG, "Attempt cleanup on unknown network, slot");
+            return;
+        }
+        cleanupAutoOnOffKeepalive(autoKi);
     }
 
     /**
@@ -414,6 +433,9 @@
         final List<AutomaticOnOffKeepalive> matches =
                 CollectionUtils.filter(mAutomaticOnOffKeepalives, it -> it.mKi.getNai() == nai);
         for (final AutomaticOnOffKeepalive ki : matches) {
+            if (ki.mAutomaticOnOffState == STATE_SUSPENDED) {
+                mKeepaliveTracker.finalizePausedKeepalive(ki.mKi, reason);
+            }
             cleanupAutoOnOffKeepalive(ki);
         }
     }
@@ -425,9 +447,14 @@
      */
     public void handleStartKeepalive(Message message) {
         final AutomaticOnOffKeepalive autoKi = (AutomaticOnOffKeepalive) message.obj;
+        final int error = mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
+        if (error != SUCCESS) {
+            mEventLog.log("Failed to start keepalive " + autoKi.mCallback + " on "
+                    + autoKi.getNetwork() + " with error " + error);
+            return;
+        }
         mEventLog.log("Start keepalive " + autoKi.mCallback + " on " + autoKi.getNetwork());
         mKeepaliveStatsTracker.onStartKeepalive();
-        mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
 
         // Add automatic on/off request into list to track its life cycle.
         try {
@@ -443,10 +470,22 @@
         }
     }
 
-    private void handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+    /**
+     * Handle resume keepalive with the given KeepaliveInfo
+     *
+     * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
+     */
+    private int handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+        final int error = mKeepaliveTracker.handleStartKeepalive(ki);
+        if (error != SUCCESS) {
+            mEventLog.log("Failed to resume keepalive " + ki.mCallback + " on " + ki.mNai
+                    + " with error " + error);
+            return error;
+        }
         mKeepaliveStatsTracker.onResumeKeepalive();
-        mKeepaliveTracker.handleStartKeepalive(ki);
         mEventLog.log("Resumed successfully keepalive " + ki.mCallback + " on " + ki.mNai);
+
+        return SUCCESS;
     }
 
     private void handlePauseKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
@@ -467,7 +506,7 @@
             final KeepaliveTracker.KeepaliveInfo ki = autoKi.mKi;
             mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), reason);
         } else {
-            mKeepaliveTracker.finalizePausedKeepalive(autoKi.mKi);
+            mKeepaliveTracker.finalizePausedKeepalive(autoKi.mKi, reason);
         }
 
         cleanupAutoOnOffKeepalive(autoKi);
@@ -612,7 +651,22 @@
      * Forward to KeepaliveTracker.
      */
     public void handleCheckKeepalivesStillValid(NetworkAgentInfo nai) {
-        mKeepaliveTracker.handleCheckKeepalivesStillValid(nai);
+        ArrayList<Pair<AutomaticOnOffKeepalive, Integer>> invalidKeepalives = null;
+
+        for (final AutomaticOnOffKeepalive autoKi : mAutomaticOnOffKeepalives) {
+            if (!nai.equals(autoKi.mKi.mNai)) continue;
+            final int error = autoKi.mKi.isValid();
+            if (error != SUCCESS) {
+                if (invalidKeepalives == null) {
+                    invalidKeepalives = new ArrayList<>();
+                }
+                invalidKeepalives.add(Pair.create(autoKi, error));
+            }
+        }
+        if (invalidKeepalives == null) return;
+        for (final Pair<AutomaticOnOffKeepalive, Integer> keepaliveAndError : invalidKeepalives) {
+            handleStopKeepalive(keepaliveAndError.first, keepaliveAndError.second);
+        }
     }
 
     @VisibleForTesting
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index cc226ce..941b616 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -54,7 +54,6 @@
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
-import android.util.Pair;
 
 import com.android.connectivity.resources.R;
 import com.android.internal.annotations.VisibleForTesting;
@@ -337,7 +336,12 @@
             return SUCCESS;
         }
 
-        private int isValid() {
+        /**
+         * Checks if the keepalive info is valid to start.
+         *
+         * @return SUCCESS if the keepalive is valid and the error reason otherwise.
+         */
+        public int isValid() {
             synchronized (mNai) {
                 int error = checkInterval();
                 if (error == SUCCESS) error = checkLimit();
@@ -348,11 +352,17 @@
             }
         }
 
-        void start(int slot) {
+        /**
+         * Attempt to start the keepalive on the given slot.
+         *
+         * @param slot the slot to start the keepalive on.
+         * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
+         */
+        int start(int slot) {
             // BINDER_DIED can happen if the binder died before the KeepaliveInfo was created and
             // the constructor set the state to BINDER_DIED. If that's the case, the KI is already
             // cleaned up.
-            if (BINDER_DIED == mStartedState) return;
+            if (BINDER_DIED == mStartedState) return BINDER_DIED;
             mSlot = slot;
             int error = isValid();
             if (error == SUCCESS) {
@@ -368,7 +378,7 @@
                             mTcpController.startSocketMonitor(mFd, this, mSlot);
                         } catch (InvalidSocketException e) {
                             handleStopKeepalive(mNai, mSlot, ERROR_INVALID_SOCKET);
-                            return;
+                            return ERROR_INVALID_SOCKET;
                         }
                         final TcpKeepalivePacketData tcpData = (TcpKeepalivePacketData) mPacket;
                         mNai.onAddTcpKeepalivePacketFilter(slot, tcpData);
@@ -377,13 +387,14 @@
                         break;
                     default:
                         Log.wtf(TAG, "Starting keepalive with unknown type: " + mType);
-                        handleStopKeepalive(mNai, mSlot, error);
-                        return;
+                        handleStopKeepalive(mNai, mSlot, ERROR_UNSUPPORTED);
+                        return ERROR_UNSUPPORTED;
                 }
                 mStartedState = STARTING;
+                return SUCCESS;
             } else {
                 handleStopKeepalive(mNai, mSlot, error);
-                return;
+                return error;
             }
         }
 
@@ -444,6 +455,8 @@
             }
         }
 
+        // TODO: This does not clean up the autoKi in AutomaticOnOffKeepaliveTracker and it is not
+        // possible without a big refactor.
         void onFileDescriptorInitiatedStop(final int socketKeepaliveReason) {
             handleStopKeepalive(mNai, mSlot, socketKeepaliveReason);
         }
@@ -486,12 +499,15 @@
 
     /**
      * Handle start keepalives with the message.
+     *
+     * @param ki the keepalive to start.
+     * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
      */
-    public void handleStartKeepalive(KeepaliveInfo ki) {
+    public int handleStartKeepalive(KeepaliveInfo ki) {
         NetworkAgentInfo nai = ki.getNai();
         int slot = findFirstFreeSlot(nai);
         mKeepalives.get(nai).put(slot, ki);
-        ki.start(slot);
+        return ki.start(slot);
     }
 
     public void handleStopAllKeepalives(NetworkAgentInfo nai, int reason) {
@@ -593,40 +609,33 @@
     /**
      * Finalize a paused keepalive.
      *
-     * This will simply send the onStopped() callback after checking that this keepalive is
-     * indeed paused.
+     * This will send the appropriate callback after checking that this keepalive is indeed paused.
      *
      * @param ki the keepalive to finalize
+     * @param reason the reason the keepalive is stopped
      */
-    public void finalizePausedKeepalive(@NonNull final KeepaliveInfo ki) {
+    public void finalizePausedKeepalive(@NonNull final KeepaliveInfo ki, int reason) {
         if (SUCCESS_PAUSED != ki.mStopReason) {
             throw new IllegalStateException("Keepalive is not paused");
         }
-        try {
-            ki.mCallback.onStopped();
-        } catch (RemoteException e) {
-            Log.w(TAG, "Discarded onStopped callback while finalizing paused keepalive");
+        if (reason == SUCCESS) {
+            try {
+                ki.mCallback.onStopped();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Discarded onStopped callback while finalizing paused keepalive");
+            }
+        } else {
+            notifyErrorCallback(ki.mCallback, reason);
         }
     }
 
-    public void handleCheckKeepalivesStillValid(NetworkAgentInfo nai) {
-        HashMap <Integer, KeepaliveInfo> networkKeepalives = mKeepalives.get(nai);
-        if (networkKeepalives != null) {
-            ArrayList<Pair<Integer, Integer>> invalidKeepalives = new ArrayList<>();
-            for (int slot : networkKeepalives.keySet()) {
-                int error = networkKeepalives.get(slot).isValid();
-                if (error != SUCCESS) {
-                    invalidKeepalives.add(Pair.create(slot, error));
-                }
-            }
-            for (Pair<Integer, Integer> slotAndError: invalidKeepalives) {
-                handleStopKeepalive(nai, slotAndError.first, slotAndError.second);
-            }
-        }
-    }
-
-    /** Handle keepalive events from lower layer. */
-    public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
+    /**
+     * Handle keepalive events from lower layer.
+     *
+     * @return false if the event caused handleStopKeepalive to be called, i.e. the keepalive is
+     *     forced to stop. Otherwise, return true.
+     */
+    public boolean handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
         KeepaliveInfo ki = null;
         try {
             ki = mKeepalives.get(nai).get(slot);
@@ -634,7 +643,7 @@
         if (ki == null) {
             Log.e(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for unknown keepalive " + slot + " on " + nai.toShortString());
-            return;
+            return true;
         }
 
         // This can be called in a number of situations :
@@ -667,11 +676,13 @@
                     Log.w(TAG, "Discarded " + (ki.mResumed ? "onResumed" : "onStarted")
                             + " callback for slot " + slot);
                 }
+                return true;
             } else {
                 Log.d(TAG, "Failed to start keepalive " + slot + " on " + nai.toShortString()
                         + ": " + reason);
                 // The message indicated some error trying to start: do call handleStopKeepalive.
                 handleStopKeepalive(nai, slot, reason);
+                return false;
             }
         } else if (KeepaliveInfo.STOPPING == ki.mStartedState) {
             // The message indicated result of stopping : clean up keepalive slots.
@@ -679,9 +690,12 @@
                     + " stopped: " + reason);
             ki.mStartedState = KeepaliveInfo.NOT_STARTED;
             cleanupStoppedKeepalive(nai, slot);
+            return true;
         } else {
             Log.wtf(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for keepalive in wrong state: " + ki.toString());
+            // Although this is an unexpected event, the keepalive is not stopped here.
+            return true;
         }
     }
 
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index 9e0435d..608e6d8 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -20,9 +20,12 @@
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 
+import static com.android.testutils.HandlerUtils.visibleOnHandlerThread;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -30,17 +33,20 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.longThat;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.ignoreStubs;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.app.AlarmManager;
 import android.content.Context;
 import android.content.res.Resources;
 import android.net.INetd;
 import android.net.ISocketKeepaliveCallback;
-import android.net.KeepalivePacketData;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.MarkMaskParcel;
@@ -48,6 +54,7 @@
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
+import android.net.SocketKeepalive;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
@@ -63,6 +70,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.connectivity.resources.R;
+import com.android.server.connectivity.AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive;
 import com.android.server.connectivity.KeepaliveTracker.KeepaliveInfo;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
@@ -70,6 +78,7 @@
 
 import libcore.util.HexEncoding;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -82,12 +91,15 @@
 import java.net.Socket;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
 public class AutomaticOnOffKeepaliveTrackerTest {
     private static final String TAG = AutomaticOnOffKeepaliveTrackerTest.class.getSimpleName();
+    private static final int TEST_SLOT = 1;
     private static final int TEST_NETID = 0xA85;
     private static final int TEST_NETID_FWMARK = 0x0A85;
     private static final int OTHER_NETID = 0x1A85;
@@ -95,6 +107,8 @@
     private static final int TIMEOUT_MS = 30_000;
     private static final int MOCK_RESOURCE_ID = 5;
     private static final int TEST_KEEPALIVE_INTERVAL_SEC = 10;
+    private static final int TEST_KEEPALIVE_INVALID_INTERVAL_SEC = 9;
+
     private AutomaticOnOffKeepaliveTracker mAOOKeepaliveTracker;
     private HandlerThread mHandlerThread;
 
@@ -102,6 +116,8 @@
     @Mock AutomaticOnOffKeepaliveTracker.Dependencies mDependencies;
     @Mock Context mCtx;
     @Mock AlarmManager mAlarmManager;
+    @Mock NetworkAgentInfo mNai;
+
     TestKeepaliveTracker mKeepaliveTracker;
     AOOTestHandler mTestHandler;
 
@@ -202,6 +218,37 @@
     private static final byte[] TEST_RESPONSE_BYTES =
             HexEncoding.decode(TEST_RESPONSE_HEX.toCharArray(), false);
 
+    private static class TestKeepaliveInfo {
+        private static List<Socket> sOpenSockets = new ArrayList<>();
+
+        public static void closeAllSockets() throws Exception {
+            for (final Socket socket : sOpenSockets) {
+                socket.close();
+            }
+            sOpenSockets.clear();
+        }
+
+        public final Socket socket;
+        public final Binder binder;
+        public final FileDescriptor fd;
+        public final ISocketKeepaliveCallback socketKeepaliveCallback;
+        public final Network underpinnedNetwork;
+        public final NattKeepalivePacketData kpd;
+
+        TestKeepaliveInfo(NattKeepalivePacketData kpd) throws Exception {
+            this.kpd = kpd;
+            socket = new Socket();
+            socket.bind(null);
+            sOpenSockets.add(socket);
+            fd = socket.getFileDescriptor$();
+
+            binder = new Binder();
+            socketKeepaliveCallback = mock(ISocketKeepaliveCallback.class);
+            doReturn(binder).when(socketKeepaliveCallback).asBinder();
+            underpinnedNetwork = mock(Network.class);
+        }
+    }
+
     private class TestKeepaliveTracker extends KeepaliveTracker {
         private KeepaliveInfo mKi;
 
@@ -231,6 +278,14 @@
     public void setup() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mNai.networkCapabilities =
+                new NetworkCapabilities.Builder().addTransportType(TRANSPORT_CELLULAR).build();
+        mNai.networkInfo = new NetworkInfo(TYPE_MOBILE, 0 /* subtype */, "LTE", "LTE");
+        mNai.networkInfo.setDetailedState(
+                NetworkInfo.DetailedState.CONNECTED, "test reason", "test extra info");
+        doReturn(new Network(TEST_NETID)).when(mNai).network();
+        mNai.linkProperties = new LinkProperties();
+
         doReturn(PERMISSION_GRANTED).when(mCtx).checkPermission(any() /* permission */,
                 anyInt() /* pid */, anyInt() /* uid */);
         ConnectivityResources.setResourcesContextForTest(mCtx);
@@ -255,6 +310,11 @@
                 new AutomaticOnOffKeepaliveTracker(mCtx, mTestHandler, mDependencies);
     }
 
+    @After
+    public void teardown() throws Exception {
+        TestKeepaliveInfo.closeAllSockets();
+    }
+
     private final class AOOTestHandler extends Handler {
         public AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive mLastAutoKi = null;
 
@@ -305,45 +365,70 @@
                 () -> assertFalse(mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
     }
 
-    @Test
-    public void testAlarm() throws Exception {
+    private void triggerEventKeepalive(int slot, int reason) {
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleEventSocketKeepalive(mNai, slot, reason));
+    }
+
+    private TestKeepaliveInfo doStartNattKeepalive(int intervalSeconds) throws Exception {
         final InetAddress srcAddress = InetAddress.getByAddress(
                 new byte[] { (byte) 192, 0, 0, (byte) 129 });
         final int srcPort = 12345;
-        final InetAddress dstAddress = InetAddress.getByAddress(new byte[] { 8, 8, 8, 8});
+        final InetAddress dstAddress = InetAddress.getByAddress(new byte[] {8, 8, 8, 8});
         final int dstPort = 12345;
 
-        final NetworkAgentInfo nai = mock(NetworkAgentInfo.class);
-        nai.networkCapabilities = new NetworkCapabilities.Builder()
-                .addTransportType(TRANSPORT_CELLULAR).build();
-        nai.networkInfo = new NetworkInfo(TYPE_MOBILE, 0 /* subtype */, "LTE", "LTE");
-        nai.networkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, "test reason",
-                "test extra info");
-        nai.linkProperties = new LinkProperties();
-        nai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
+        mNai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
 
-        final Socket socket = new Socket();
-        socket.bind(null);
-        final FileDescriptor fd = socket.getFileDescriptor$();
-        final IBinder binder = new Binder();
-        final ISocketKeepaliveCallback cb = mock(ISocketKeepaliveCallback.class);
-        doReturn(binder).when(cb).asBinder();
-        final Network underpinnedNetwork = mock(Network.class);
-
-        final KeepalivePacketData kpd = new NattKeepalivePacketData(srcAddress, srcPort,
+        final NattKeepalivePacketData kpd = new NattKeepalivePacketData(srcAddress, srcPort,
                 dstAddress, dstPort, new byte[] {1});
-        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(cb, nai, kpd,
-                TEST_KEEPALIVE_INTERVAL_SEC, KeepaliveInfo.TYPE_NATT, fd);
+
+        final TestKeepaliveInfo testInfo = new TestKeepaliveInfo(kpd);
+
+        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(
+                testInfo.socketKeepaliveCallback, mNai, kpd, intervalSeconds,
+                KeepaliveInfo.TYPE_NATT, testInfo.fd);
         mKeepaliveTracker.setReturnedKeepaliveInfo(ki);
 
+        mAOOKeepaliveTracker.startNattKeepalive(mNai, testInfo.fd, intervalSeconds,
+                testInfo.socketKeepaliveCallback, srcAddress.toString(), srcPort,
+                dstAddress.toString(), dstPort, true /* automaticOnOffKeepalives */,
+                testInfo.underpinnedNetwork);
+        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+
+        return testInfo;
+    }
+
+    private TestKeepaliveInfo doStartNattKeepalive() throws Exception {
+        return doStartNattKeepalive(TEST_KEEPALIVE_INTERVAL_SEC);
+    }
+
+    private void doPauseKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        setupResponseWithoutSocketExisting();
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleMonitorAutomaticKeepalive(autoKi, TEST_NETID));
+    }
+
+    private void doResumeKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        setupResponseWithSocketExisting();
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleMonitorAutomaticKeepalive(autoKi, TEST_NETID));
+    }
+
+    private void doStopKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleStopKeepalive(autoKi, SocketKeepalive.SUCCESS));
+    }
+
+    @Test
+    public void testAlarm() throws Exception {
         // Mock elapsed real time to verify the alarm timer.
         final long time = SystemClock.elapsedRealtime();
         doReturn(time).when(mDependencies).getElapsedRealtime();
-
-        mAOOKeepaliveTracker.startNattKeepalive(nai, fd, 10 /* intervalSeconds */, cb,
-                srcAddress.toString(), srcPort, dstAddress.toString(), dstPort,
-                true /* automaticOnOffKeepalives */, underpinnedNetwork);
-        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
 
         final ArgumentCaptor<AlarmManager.OnAlarmListener> listenerCaptor =
                 ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
@@ -362,9 +447,8 @@
         HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
 
         assertNotNull(mTestHandler.mLastAutoKi);
-        assertEquals(cb, mTestHandler.mLastAutoKi.getCallback());
-        assertEquals(underpinnedNetwork, mTestHandler.mLastAutoKi.getUnderpinnedNetwork());
-        socket.close();
+        assertEquals(testInfo.socketKeepaliveCallback, mTestHandler.mLastAutoKi.getCallback());
+        assertEquals(testInfo.underpinnedNetwork, mTestHandler.mLastAutoKi.getUnderpinnedNetwork());
     }
 
     private void setupResponseWithSocketExisting() throws Exception {
@@ -391,4 +475,301 @@
         buffer.order(ByteOrder.nativeOrder());
         return buffer;
     }
+
+    private AutomaticOnOffKeepalive getAutoKiForBinder(IBinder binder) {
+        return visibleOnHandlerThread(
+                mTestHandler, () -> mAOOKeepaliveTracker.getKeepaliveForBinder(binder));
+    }
+
+    private void checkAndProcessKeepaliveStart(final NattKeepalivePacketData kpd) throws Exception {
+        checkAndProcessKeepaliveStart(TEST_SLOT, kpd);
+    }
+
+    private void checkAndProcessKeepaliveStart(
+            int slot, final NattKeepalivePacketData kpd) throws Exception {
+        verify(mNai).onStartNattSocketKeepalive(slot, TEST_KEEPALIVE_INTERVAL_SEC, kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(slot, kpd);
+        triggerEventKeepalive(slot, SocketKeepalive.SUCCESS);
+    }
+
+    private void checkAndProcessKeepaliveStop() throws Exception {
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+    }
+
+    private void checkAndProcessKeepaliveStop(int slot) throws Exception {
+        verify(mNai).onStopSocketKeepalive(slot);
+        verify(mNai).onRemoveKeepalivePacketFilter(slot);
+        triggerEventKeepalive(slot, SocketKeepalive.SUCCESS);
+    }
+
+    @Test
+    public void testStartNattKeepalive_valid() throws Exception {
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+
+        final AutomaticOnOffKeepalive autoKi = getAutoKiForBinder(testInfo.binder);
+        assertNotNull(autoKi);
+        assertEquals(testInfo.socketKeepaliveCallback, autoKi.getCallback());
+
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStartNattKeepalive_invalidInterval() throws Exception {
+        final TestKeepaliveInfo testInfo =
+                doStartNattKeepalive(TEST_KEEPALIVE_INVALID_INTERVAL_SEC);
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_INTERVAL);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testHandleEventSocketKeepalive_startingFailureHardwareError() throws Exception {
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+
+        verify(mNai)
+                .onStartNattSocketKeepalive(TEST_SLOT, TEST_KEEPALIVE_INTERVAL_SEC, testInfo.kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(TEST_SLOT, testInfo.kpd);
+        // Network agent returns an error, fails to start the keepalive.
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.ERROR_HARDWARE_ERROR);
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_HARDWARE_ERROR);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testHandleCheckKeepalivesStillValid_linkPropertiesChanged() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        // Source address is removed from link properties by clearing.
+        mNai.linkProperties.clear();
+
+        // Check for valid keepalives
+        visibleOnHandlerThread(
+                mTestHandler, () -> mAOOKeepaliveTracker.handleCheckKeepalivesStillValid(mNai));
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStopKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testPauseKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        // Pausing does not cleanup the autoKi
+        assertNotNull(getAutoKiForBinder(testInfo.binder));
+
+        clearInvocations(mNai);
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        // The keepalive is already stopped.
+        verify(mNai, never()).onStopSocketKeepalive(TEST_SLOT);
+        verify(mNai, never()).onRemoveKeepalivePacketFilter(TEST_SLOT);
+
+        // Stopping while paused still calls onStopped.
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        // autoKi is cleaned up.
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        assertNotNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onResumed();
+
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive_invalidSourceAddress() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        mNai.linkProperties.clear();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+        verify(mNai, never()).onStartNattSocketKeepalive(anyInt(), anyInt(), any());
+        verify(mNai, never()).onAddNattKeepalivePacketFilter(anyInt(), any());
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive_startingFailureHardwareError() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+
+        verify(mNai)
+                .onStartNattSocketKeepalive(TEST_SLOT, TEST_KEEPALIVE_INTERVAL_SEC, testInfo.kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(TEST_SLOT, testInfo.kpd);
+        // Network agent returns error on starting the keepalive.
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.ERROR_HARDWARE_ERROR);
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_HARDWARE_ERROR);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStopAllKeepalives() throws Exception {
+        final TestKeepaliveInfo testInfo1 = doStartNattKeepalive();
+        final TestKeepaliveInfo testInfo2 = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(TEST_SLOT, testInfo1.kpd);
+        checkAndProcessKeepaliveStart(TEST_SLOT + 1, testInfo2.kpd);
+
+        verify(testInfo1.socketKeepaliveCallback).onStarted();
+        verify(testInfo2.socketKeepaliveCallback).onStarted();
+
+        // Pause the first keepalive
+        doPauseKeepalive(getAutoKiForBinder(testInfo1.binder));
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        verify(testInfo1.socketKeepaliveCallback).onPaused();
+
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleStopAllKeepalives(
+                        mNai, SocketKeepalive.ERROR_INVALID_NETWORK));
+
+        // Note that checkAndProcessKeepaliveStop is not called since the network agent is assumed
+        // to be disconnected for a handleStopAllKeepalives call.
+        assertNull(getAutoKiForBinder(testInfo1.binder));
+        assertNull(getAutoKiForBinder(testInfo2.binder));
+
+        verify(testInfo1.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo2.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo1.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_NETWORK);
+        verify(testInfo2.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_NETWORK);
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo1.socketKeepaliveCallback));
+        verifyNoMoreInteractions(ignoreStubs(testInfo2.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testTwoKeepalives_startAfterPause() throws Exception {
+        final TestKeepaliveInfo testInfo1 = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo1.kpd);
+        verify(testInfo1.socketKeepaliveCallback).onStarted();
+        assertNotNull(getAutoKiForBinder(testInfo1.binder));
+
+        final AutomaticOnOffKeepalive autoKi1  = getAutoKiForBinder(testInfo1.binder);
+        doPauseKeepalive(autoKi1);
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        verify(testInfo1.socketKeepaliveCallback).onPaused();
+        assertNotNull(getAutoKiForBinder(testInfo1.binder));
+
+        clearInvocations(mNai);
+        // Start the second keepalive while the first is paused.
+        final TestKeepaliveInfo testInfo2 = doStartNattKeepalive();
+        // The slot used is TEST_SLOT since it is now a free slot.
+        checkAndProcessKeepaliveStart(TEST_SLOT, testInfo2.kpd);
+        verify(testInfo2.socketKeepaliveCallback).onStarted();
+        assertNotNull(getAutoKiForBinder(testInfo2.binder));
+
+        clearInvocations(mNai);
+        doResumeKeepalive(autoKi1);
+        // The next free slot is TEST_SLOT + 1.
+        checkAndProcessKeepaliveStart(TEST_SLOT + 1, testInfo1.kpd);
+        verify(testInfo1.socketKeepaliveCallback).onResumed();
+
+        clearInvocations(mNai);
+        doStopKeepalive(autoKi1);
+        // TODO: The slot should be consistent with the checkAndProcessKeepaliveStart directly above
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        // TODO: onStopped should only be called on the first keepalive callback.
+        verify(testInfo1.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo2.socketKeepaliveCallback).onStopped();
+        assertNull(getAutoKiForBinder(testInfo1.binder));
+
+        clearInvocations(mNai);
+        assertNotNull(getAutoKiForBinder(testInfo2.binder));
+        doStopKeepalive(getAutoKiForBinder(testInfo2.binder));
+        // This slot should be consistent with its corresponding checkAndProcessKeepaliveStart.
+        // TODO: checkAndProcessKeepaliveStop should be called instead but the keepalive is
+        // unexpectedly already stopped above.
+        verify(mNai, never()).onStopSocketKeepalive(TEST_SLOT);
+        verify(mNai, never()).onRemoveKeepalivePacketFilter(TEST_SLOT);
+
+        verify(testInfo2.socketKeepaliveCallback).onStopped();
+        assertNull(getAutoKiForBinder(testInfo2.binder));
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo1.socketKeepaliveCallback));
+        verifyNoMoreInteractions(ignoreStubs(testInfo2.socketKeepaliveCallback));
+    }
 }