diff --git a/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java b/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
index 8c5d1d6..8b072c4 100644
--- a/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
+++ b/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
@@ -22,7 +22,9 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
@@ -56,6 +58,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -80,6 +83,12 @@
 
     IpConnectivityMetrics mService;
     NetdEventListenerService mNetdListener;
+    private static final NetworkCapabilities CAPABILITIES_WIFI = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    private static final NetworkCapabilities CAPABILITIES_CELL = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
@@ -263,14 +272,6 @@
         // TODO: instead of comparing textpb to textpb, parse textpb and compare proto to proto.
         IpConnectivityLog logger = new IpConnectivityLog(mService.impl);
 
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         ApfStats apfStats = new ApfStats.Builder()
                 .setDurationMs(45000)
                 .setReceivedRas(10)
@@ -584,11 +585,21 @@
         return buffer.toString();
     }
 
-    void connectEvent(int netid, int error, int latencyMs, String ipAddr) throws Exception {
-        mNetdListener.onConnectEvent(netid, error, latencyMs, ipAddr, 80, 1);
+    private void setCapabilities(int netId) {
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? CAPABILITIES_WIFI : CAPABILITIES_CELL);
+    }
+
+    void connectEvent(int netId, int error, int latencyMs, String ipAddr) throws Exception {
+        setCapabilities(netId);
+        mNetdListener.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mNetdListener.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }
 
diff --git a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
index 8ccea1a..50aaaee 100644
--- a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
+++ b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
@@ -23,8 +23,9 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.net.ConnectivityManager;
@@ -42,6 +43,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.io.FileOutputStream;
 import java.io.PrintWriter;
@@ -61,18 +63,16 @@
 
     NetdEventListenerService mService;
     ConnectivityManager mCm;
+    private static final NetworkCapabilities CAPABILITIES_WIFI = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    private static final NetworkCapabilities CAPABILITIES_CELL = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
         mCm = mock(ConnectivityManager.class);
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         mService = new NetdEventListenerService(mCm);
     }
 
@@ -470,7 +470,16 @@
         assertEquals(want, got);
     }
 
+    private void setCapabilities(int netId) {
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? CAPABILITIES_WIFI : CAPABILITIES_CELL);
+    }
+
     Thread connectEventAction(int netId, int error, int latencyMs, String ipAddr) {
+        setCapabilities(netId);
         return new Thread(() -> {
             try {
                 mService.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
@@ -481,6 +490,7 @@
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mService.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }