diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 0e10de8..a69ca99 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -3449,6 +3449,10 @@
             final NetworkCallback callback;
             synchronized (sCallbacks) {
                 callback = sCallbacks.get(request);
+                if (message.what == CALLBACK_UNAVAIL) {
+                    sCallbacks.remove(request);
+                    callback.networkRequest = ALREADY_UNREGISTERED;
+                }
             }
             if (DBG) {
                 Log.d(TAG, getCallbackName(message.what) + " for network " + network);
@@ -3995,8 +3999,10 @@
         synchronized (sCallbacks) {
             Preconditions.checkArgument(networkCallback.networkRequest != null,
                     "NetworkCallback was not registered");
-            Preconditions.checkArgument(networkCallback.networkRequest != ALREADY_UNREGISTERED,
-                    "NetworkCallback was already unregistered");
+            if (networkCallback.networkRequest == ALREADY_UNREGISTERED) {
+                Log.d(TAG, "NetworkCallback was already unregistered");
+                return;
+            }
             for (Map.Entry<NetworkRequest, NetworkCallback> e : sCallbacks.entrySet()) {
                 if (e.getValue() == networkCallback) {
                     reqs.add(e.getKey());
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index ae58e99..40aec8f 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -462,10 +462,11 @@
         if (networkKeepalives != null) {
             for (KeepaliveInfo ki : networkKeepalives.values()) {
                 ki.stop(reason);
+                // Clean up keepalives since the network agent is disconnected and unable to pass
+                // back asynchronous result of stop().
+                cleanupStoppedKeepalive(nai, ki.mSlot);
             }
         }
-        // Clean up keepalives will be done as a result of calling ki.stop() after the slots are
-        // freed.
     }
 
     public void handleStopKeepalive(NetworkAgentInfo nai, int slot, int reason) {
diff --git a/tests/net/Android.bp b/tests/net/Android.bp
index 1fbb658..3fe867b 100644
--- a/tests/net/Android.bp
+++ b/tests/net/Android.bp
@@ -29,7 +29,6 @@
         "libbpf",
         "libbpf_android",
         "libc++",
-        "libcgrouprc",
         "libcrypto",
         "libcutils",
         "libdexfile",
diff --git a/tests/net/java/android/net/IpMemoryStoreTest.java b/tests/net/java/android/net/IpMemoryStoreTest.java
index 18c6768..8ff2de9 100644
--- a/tests/net/java/android/net/IpMemoryStoreTest.java
+++ b/tests/net/java/android/net/IpMemoryStoreTest.java
@@ -16,10 +16,26 @@
 
 package android.net;
 
-import static org.mockito.ArgumentMatchers.any;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.content.Context;
+import android.net.ipmemorystore.Blob;
+import android.net.ipmemorystore.IOnStatusListener;
+import android.net.ipmemorystore.NetworkAttributes;
+import android.net.ipmemorystore.NetworkAttributesParcelable;
+import android.net.ipmemorystore.Status;
+import android.os.RemoteException;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -27,28 +43,57 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.net.UnknownHostException;
+import java.util.Arrays;
+
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class IpMemoryStoreTest {
+    private static final String TAG = IpMemoryStoreTest.class.getSimpleName();
+    private static final String TEST_CLIENT_ID = "testClientId";
+    private static final String TEST_DATA_NAME = "testData";
+    private static final String TEST_OTHER_DATA_NAME = TEST_DATA_NAME + "Other";
+    private static final byte[] TEST_BLOB_DATA = new byte[] { -3, 6, 8, -9, 12,
+            -128, 0, 89, 112, 91, -34 };
+    private static final NetworkAttributes TEST_NETWORK_ATTRIBUTES = buildTestNetworkAttributes(
+            "hint", 219);
+
     @Mock
     Context mMockContext;
     @Mock
     NetworkStackClient mNetworkStackClient;
     @Mock
     IIpMemoryStore mMockService;
+    @Mock
+    IOnStatusListener mIOnStatusListener;
     IpMemoryStore mStore;
 
+    @Captor
+    ArgumentCaptor<IIpMemoryStoreCallbacks> mCbCaptor;
+    @Captor
+    ArgumentCaptor<NetworkAttributesParcelable> mNapCaptor;
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        doAnswer(invocation -> {
-            ((IIpMemoryStoreCallbacks) invocation.getArgument(0))
-                    .onIpMemoryStoreFetched(mMockService);
-            return null;
-        }).when(mNetworkStackClient).fetchIpMemoryStore(any());
+    }
+
+    private void startIpMemoryStore(boolean supplyService) {
+        if (supplyService) {
+            doAnswer(invocation -> {
+                ((IIpMemoryStoreCallbacks) invocation.getArgument(0))
+                        .onIpMemoryStoreFetched(mMockService);
+                return null;
+            }).when(mNetworkStackClient).fetchIpMemoryStore(any());
+        } else {
+            doNothing().when(mNetworkStackClient).fetchIpMemoryStore(mCbCaptor.capture());
+        }
         mStore = new IpMemoryStore(mMockContext) {
             @Override
             protected NetworkStackClient getNetworkStackClient() {
@@ -57,24 +102,228 @@
         };
     }
 
-    @Test
-    public void testNetworkAttributes() {
-        // TODO : implement this
+    private static NetworkAttributes buildTestNetworkAttributes(String hint, int mtu) {
+        return new NetworkAttributes.Builder()
+                .setGroupHint(hint)
+                .setMtu(mtu)
+                .build();
     }
 
     @Test
-    public void testPrivateData() {
-        // TODO : implement this
+    public void testNetworkAttributes() throws Exception {
+        startIpMemoryStore(true);
+        final String l2Key = "fakeKey";
+
+        mStore.storeNetworkAttributes(l2Key, TEST_NETWORK_ATTRIBUTES,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        verify(mMockService, times(1)).storeNetworkAttributes(eq(l2Key),
+                mNapCaptor.capture(), any());
+        assertEquals(TEST_NETWORK_ATTRIBUTES, new NetworkAttributes(mNapCaptor.getValue()));
+
+        mStore.retrieveNetworkAttributes(l2Key,
+                (status, key, attr) -> {
+                    assertTrue("Retrieve network attributes not successful : "
+                            + status.resultCode, status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(TEST_NETWORK_ATTRIBUTES, attr);
+                });
+
+        verify(mMockService, times(1)).retrieveNetworkAttributes(eq(l2Key), any());
     }
 
     @Test
-    public void testFindL2Key() {
-        // TODO : implement this
+    public void testPrivateData() throws RemoteException {
+        startIpMemoryStore(true);
+        final Blob b = new Blob();
+        b.data = TEST_BLOB_DATA;
+        final String l2Key = "fakeKey";
+
+        mStore.storeBlob(l2Key, TEST_CLIENT_ID, TEST_DATA_NAME, b,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        verify(mMockService, times(1)).storeBlob(eq(l2Key), eq(TEST_CLIENT_ID), eq(TEST_DATA_NAME),
+                eq(b), any());
+
+        mStore.retrieveBlob(l2Key, TEST_CLIENT_ID, TEST_OTHER_DATA_NAME,
+                (status, key, name, data) -> {
+                    assertTrue("Retrieve blob status not successful : " + status.resultCode,
+                            status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(name, TEST_DATA_NAME);
+                    assertTrue(Arrays.equals(b.data, data.data));
+                });
+        verify(mMockService, times(1)).retrieveBlob(eq(l2Key), eq(TEST_CLIENT_ID),
+                eq(TEST_OTHER_DATA_NAME), any());
     }
 
     @Test
-    public void testIsSameNetwork() {
-        // TODO : implement this
+    public void testFindL2Key()
+            throws UnknownHostException, RemoteException, Exception {
+        startIpMemoryStore(true);
+        final String l2Key = "fakeKey";
+
+        mStore.findL2Key(TEST_NETWORK_ATTRIBUTES,
+                (status, key) -> {
+                    assertTrue("Retrieve network sameness not successful : " + status.resultCode,
+                            status.isSuccess());
+                    assertEquals(l2Key, key);
+                });
+        verify(mMockService, times(1)).findL2Key(mNapCaptor.capture(), any());
+        assertEquals(TEST_NETWORK_ATTRIBUTES, new NetworkAttributes(mNapCaptor.getValue()));
     }
 
+    @Test
+    public void testIsSameNetwork() throws UnknownHostException, RemoteException {
+        startIpMemoryStore(true);
+        final String l2Key1 = "fakeKey1";
+        final String l2Key2 = "fakeKey2";
+
+        mStore.isSameNetwork(l2Key1, l2Key2,
+                (status, answer) -> {
+                    assertFalse("Retrieve network sameness suspiciously successful : "
+                            + status.resultCode, status.isSuccess());
+                    assertEquals(Status.ERROR_ILLEGAL_ARGUMENT, status.resultCode);
+                    assertNull(answer);
+                });
+        verify(mMockService, times(1)).isSameNetwork(eq(l2Key1), eq(l2Key2), any());
+    }
+
+    @Test
+    public void testEnqueuedIpMsRequests() throws Exception {
+        startIpMemoryStore(false);
+
+        final Blob b = new Blob();
+        b.data = TEST_BLOB_DATA;
+        final String l2Key = "fakeKey";
+
+        // enqueue multiple ipms requests
+        mStore.storeNetworkAttributes(l2Key, TEST_NETWORK_ATTRIBUTES,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        mStore.retrieveNetworkAttributes(l2Key,
+                (status, key, attr) -> {
+                    assertTrue("Retrieve network attributes not successful : "
+                            + status.resultCode, status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(TEST_NETWORK_ATTRIBUTES, attr);
+                });
+        mStore.storeBlob(l2Key, TEST_CLIENT_ID, TEST_DATA_NAME, b,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        mStore.retrieveBlob(l2Key, TEST_CLIENT_ID, TEST_OTHER_DATA_NAME,
+                (status, key, name, data) -> {
+                    assertTrue("Retrieve blob status not successful : " + status.resultCode,
+                            status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(name, TEST_DATA_NAME);
+                    assertTrue(Arrays.equals(b.data, data.data));
+                });
+
+        // get ipms service ready
+        mCbCaptor.getValue().onIpMemoryStoreFetched(mMockService);
+
+        InOrder inOrder = inOrder(mMockService);
+
+        inOrder.verify(mMockService).storeNetworkAttributes(eq(l2Key), mNapCaptor.capture(), any());
+        inOrder.verify(mMockService).retrieveNetworkAttributes(eq(l2Key), any());
+        inOrder.verify(mMockService).storeBlob(eq(l2Key), eq(TEST_CLIENT_ID), eq(TEST_DATA_NAME),
+                eq(b), any());
+        inOrder.verify(mMockService).retrieveBlob(eq(l2Key), eq(TEST_CLIENT_ID),
+                eq(TEST_OTHER_DATA_NAME), any());
+        assertEquals(TEST_NETWORK_ATTRIBUTES, new NetworkAttributes(mNapCaptor.getValue()));
+    }
+
+    @Test
+    public void testEnqueuedIpMsRequestsWithException() throws Exception {
+        startIpMemoryStore(true);
+        doThrow(RemoteException.class).when(mMockService).retrieveNetworkAttributes(any(), any());
+
+        final Blob b = new Blob();
+        b.data = TEST_BLOB_DATA;
+        final String l2Key = "fakeKey";
+
+        // enqueue multiple ipms requests
+        mStore.storeNetworkAttributes(l2Key, TEST_NETWORK_ATTRIBUTES,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        mStore.retrieveNetworkAttributes(l2Key,
+                (status, key, attr) -> {
+                    assertTrue("Retrieve network attributes not successful : "
+                            + status.resultCode, status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(TEST_NETWORK_ATTRIBUTES, attr);
+                });
+        mStore.storeBlob(l2Key, TEST_CLIENT_ID, TEST_DATA_NAME, b,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        mStore.retrieveBlob(l2Key, TEST_CLIENT_ID, TEST_OTHER_DATA_NAME,
+                (status, key, name, data) -> {
+                    assertTrue("Retrieve blob status not successful : " + status.resultCode,
+                            status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(name, TEST_DATA_NAME);
+                    assertTrue(Arrays.equals(b.data, data.data));
+                });
+
+        // verify the rest of the queue is still processed in order even if the remote exception
+        // occurs when calling one or more requests
+        InOrder inOrder = inOrder(mMockService);
+
+        inOrder.verify(mMockService).storeNetworkAttributes(eq(l2Key), mNapCaptor.capture(), any());
+        inOrder.verify(mMockService).storeBlob(eq(l2Key), eq(TEST_CLIENT_ID), eq(TEST_DATA_NAME),
+                eq(b), any());
+        inOrder.verify(mMockService).retrieveBlob(eq(l2Key), eq(TEST_CLIENT_ID),
+                eq(TEST_OTHER_DATA_NAME), any());
+        assertEquals(TEST_NETWORK_ATTRIBUTES, new NetworkAttributes(mNapCaptor.getValue()));
+    }
+
+    @Test
+    public void testEnqueuedIpMsRequestsCallbackFunctionWithException() throws Exception {
+        startIpMemoryStore(true);
+
+        final Blob b = new Blob();
+        b.data = TEST_BLOB_DATA;
+        final String l2Key = "fakeKey";
+
+        // enqueue multiple ipms requests
+        mStore.storeNetworkAttributes(l2Key, TEST_NETWORK_ATTRIBUTES,
+                status -> {
+                    assertTrue("Store not successful : " + status.resultCode, status.isSuccess());
+                });
+        mStore.retrieveNetworkAttributes(l2Key,
+                (status, key, attr) -> {
+                    throw new RuntimeException("retrieveNetworkAttributes test");
+                });
+        mStore.storeBlob(l2Key, TEST_CLIENT_ID, TEST_DATA_NAME, b,
+                status -> {
+                    throw new RuntimeException("storeBlob test");
+                });
+        mStore.retrieveBlob(l2Key, TEST_CLIENT_ID, TEST_OTHER_DATA_NAME,
+                (status, key, name, data) -> {
+                    assertTrue("Retrieve blob status not successful : " + status.resultCode,
+                            status.isSuccess());
+                    assertEquals(l2Key, key);
+                    assertEquals(name, TEST_DATA_NAME);
+                    assertTrue(Arrays.equals(b.data, data.data));
+                });
+
+        // verify the rest of the queue is still processed in order even if when one or more
+        // callback throw the remote exception
+        InOrder inOrder = inOrder(mMockService);
+
+        inOrder.verify(mMockService).storeNetworkAttributes(eq(l2Key), mNapCaptor.capture(),
+                any());
+        inOrder.verify(mMockService).storeBlob(eq(l2Key), eq(TEST_CLIENT_ID), eq(TEST_DATA_NAME),
+                eq(b), any());
+        inOrder.verify(mMockService).retrieveBlob(eq(l2Key), eq(TEST_CLIENT_ID),
+                eq(TEST_OTHER_DATA_NAME), any());
+        assertEquals(TEST_NETWORK_ATTRIBUTES, new NetworkAttributes(mNapCaptor.getValue()));
+    }
 }
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index c5bfc4b..3a28aca 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -3854,6 +3854,9 @@
         networkCallback.expectCallback(CallbackState.UNAVAILABLE, null);
         testFactory.waitForRequests();
 
+        // unregister network callback - a no-op, but should not fail
+        mCm.unregisterNetworkCallback(networkCallback);
+
         testFactory.unregister();
         handlerThread.quit();
     }
diff --git a/tests/net/java/com/android/server/IpSecServiceTest.java b/tests/net/java/com/android/server/IpSecServiceTest.java
index 6b5a220..4a35015 100644
--- a/tests/net/java/com/android/server/IpSecServiceTest.java
+++ b/tests/net/java/com/android/server/IpSecServiceTest.java
@@ -118,7 +118,6 @@
     INetd mMockNetd;
     IpSecService.IpSecServiceConfiguration mMockIpSecSrvConfig;
     IpSecService mIpSecService;
-    int mUid = Os.getuid();
 
     @Before
     public void setUp() throws Exception {
@@ -666,99 +665,4 @@
         mIpSecService.releaseNetId(releasedNetId);
         assertEquals(releasedNetId, mIpSecService.reserveNetId());
     }
-
-    @Test
-    public void testLockEncapSocketForNattKeepalive() throws Exception {
-        IpSecUdpEncapResponse udpEncapResp =
-                mIpSecService.openUdpEncapsulationSocket(0, new Binder());
-        assertNotNull(udpEncapResp);
-        assertEquals(IpSecManager.Status.OK, udpEncapResp.status);
-
-        // Verify no NATT keepalive records upon startup
-        IpSecService.UserRecord userRecord = mIpSecService.mUserResourceTracker.getUserRecord(mUid);
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-
-        int nattKeepaliveResourceId =
-                mIpSecService.lockEncapSocketForNattKeepalive(udpEncapResp.resourceId, mUid);
-
-        // Validate response, and record was added
-        assertNotEquals(IpSecManager.INVALID_RESOURCE_ID, nattKeepaliveResourceId);
-        assertEquals(1, userRecord.mNattKeepaliveRecords.size());
-
-        // Validate keepalive can be released and removed.
-        mIpSecService.releaseNattKeepalive(nattKeepaliveResourceId, mUid);
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-
-        mIpSecService.closeUdpEncapsulationSocket(udpEncapResp.resourceId);
-    }
-
-    @Test
-    public void testLockEncapSocketForNattKeepaliveInvalidUid() throws Exception {
-        IpSecUdpEncapResponse udpEncapResp =
-                mIpSecService.openUdpEncapsulationSocket(0, new Binder());
-        assertNotNull(udpEncapResp);
-        assertEquals(IpSecManager.Status.OK, udpEncapResp.status);
-
-        // Verify no NATT keepalive records upon startup
-        IpSecService.UserRecord userRecord = mIpSecService.mUserResourceTracker.getUserRecord(mUid);
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-
-        try {
-            int nattKeepaliveResourceId =
-                    mIpSecService.lockEncapSocketForNattKeepalive(
-                            udpEncapResp.resourceId, mUid + 1);
-            fail("Expected SecurityException for invalid user");
-        } catch (SecurityException expected) {
-        }
-
-        // Validate keepalive was not added to lists
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-    }
-
-    @Test
-    public void testLockEncapSocketForNattKeepaliveInvalidResourceId() throws Exception {
-        // Verify no NATT keepalive records upon startup
-        IpSecService.UserRecord userRecord = mIpSecService.mUserResourceTracker.getUserRecord(mUid);
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-
-        try {
-            int nattKeepaliveResourceId =
-                    mIpSecService.lockEncapSocketForNattKeepalive(12345, mUid);
-            fail("Expected IllegalArgumentException for invalid resource ID");
-        } catch (IllegalArgumentException expected) {
-        }
-
-        // Validate keepalive was not added to lists
-        assertEquals(0, userRecord.mNattKeepaliveRecords.size());
-    }
-
-    @Test
-    public void testEncapSocketReleasedBeforeKeepaliveReleased() throws Exception {
-        IpSecUdpEncapResponse udpEncapResp =
-                mIpSecService.openUdpEncapsulationSocket(0, new Binder());
-        assertNotNull(udpEncapResp);
-        assertEquals(IpSecManager.Status.OK, udpEncapResp.status);
-
-        // Get encap socket record, verify initial starting refcount.
-        IpSecService.UserRecord userRecord = mIpSecService.mUserResourceTracker.getUserRecord(mUid);
-        IpSecService.RefcountedResource encapSocketRefcountedRecord =
-                userRecord.mEncapSocketRecords.getRefcountedResourceOrThrow(
-                        udpEncapResp.resourceId);
-        assertEquals(1, encapSocketRefcountedRecord.mRefCount);
-
-        // Verify that the reference was added
-        int nattKeepaliveResourceId =
-                mIpSecService.lockEncapSocketForNattKeepalive(udpEncapResp.resourceId, mUid);
-        assertNotEquals(IpSecManager.INVALID_RESOURCE_ID, nattKeepaliveResourceId);
-        assertEquals(2, encapSocketRefcountedRecord.mRefCount);
-
-        // Close UDP encap socket, but expect the refcountedRecord to still have a reference.
-        mIpSecService.closeUdpEncapsulationSocket(udpEncapResp.resourceId);
-        assertEquals(1, encapSocketRefcountedRecord.mRefCount);
-
-        // Verify UDP encap socket cleaned up once reference is removed. Expect -1 if cleanup
-        // was properly completed.
-        mIpSecService.releaseNattKeepalive(nattKeepaliveResourceId, mUid);
-        assertEquals(-1, encapSocketRefcountedRecord.mRefCount);
-    }
 }
diff --git a/tests/net/java/com/android/server/net/ipmemorystore/NetworkAttributesTest.java b/tests/net/java/com/android/server/net/ipmemorystore/NetworkAttributesTest.java
index a83faf3..fb84611 100644
--- a/tests/net/java/com/android/server/net/ipmemorystore/NetworkAttributesTest.java
+++ b/tests/net/java/com/android/server/net/ipmemorystore/NetworkAttributesTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.connectivity.ipmemorystore;
+package com.android.server.net.ipmemorystore;
 
 import static org.junit.Assert.assertEquals;