diff --git a/service/Android.bp b/service/Android.bp
index 3ff7a7c..02717f7 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -68,7 +68,6 @@
         "dnsresolver_aidl_interface-V9-java",
         "modules-utils-build",
         "modules-utils-shell-command-handler",
-        "modules-utils-statemachine",
         "net-utils-device-common",
         "net-utils-device-common-netlink",
         "net-utils-framework-common",
diff --git a/service/jarjar-rules.txt b/service/jarjar-rules.txt
index 4ba6837..e5d1a88 100644
--- a/service/jarjar-rules.txt
+++ b/service/jarjar-rules.txt
@@ -8,10 +8,7 @@
 # the one in com.android.internal.util
 rule android.util.IndentingPrintWriter* com.android.connectivity.@0
 rule com.android.internal.util.IndentingPrintWriter* com.android.connectivity.@0
-rule com.android.internal.util.IState* com.android.connectivity.@0
 rule com.android.internal.util.MessageUtils* com.android.connectivity.@0
-rule com.android.internal.util.State* com.android.connectivity.@0
-rule com.android.internal.util.StateMachine* com.android.connectivity.@0
 rule com.android.internal.util.WakeupMessage* com.android.connectivity.@0
 
 rule com.android.internal.messages.** com.android.connectivity.@0
diff --git a/tests/cts/net/Android.bp b/tests/cts/net/Android.bp
index 1872327..81c30b1 100644
--- a/tests/cts/net/Android.bp
+++ b/tests/cts/net/Android.bp
@@ -69,7 +69,7 @@
 // devices.
 android_test {
     name: "CtsNetTestCases",
-    defaults: ["CtsNetTestCasesDefaults"],
+    defaults: ["CtsNetTestCasesDefaults", "NetworkStackNextEnableDefaults"],
     // TODO: CTS should not depend on the entirety of the networkstack code.
     static_libs: [
         "NetworkStackApiCurrentLib",
diff --git a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
index 9f98e3f..a378aa7 100644
--- a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
+++ b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
@@ -40,7 +40,6 @@
 import android.net.cts.util.CtsNetUtils
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTPS_URL
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTP_URL
-import android.net.wifi.WifiManager
 import android.os.Build
 import android.platform.test.annotations.AppModeFull
 import android.provider.DeviceConfig
@@ -77,7 +76,7 @@
 private const val LOCALHOST_HOSTNAME = "localhost"
 
 // Re-connecting to the AP, obtaining an IP address, revalidating can take a long time
-private const val WIFI_CONNECT_TIMEOUT_MS = 120_000L
+private const val WIFI_CONNECT_TIMEOUT_MS = 40_000L
 private const val TEST_TIMEOUT_MS = 10_000L
 
 private const val TAG = "CaptivePortalTest"
@@ -94,7 +93,6 @@
 @RunWith(AndroidJUnit4::class)
 class CaptivePortalTest {
     private val context: android.content.Context by lazy { getInstrumentation().context }
-    private val wm by lazy { context.getSystemService(WifiManager::class.java) }
     private val cm by lazy { context.getSystemService(ConnectivityManager::class.java) }
     private val pm by lazy { context.packageManager }
     private val utils by lazy { CtsNetUtils(context) }
diff --git a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
index e17200e..05e40f9 100644
--- a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -306,14 +306,18 @@
         }
 
         try {
+            if (wasWifiConnected) {
+                // Make sure the callback is registered before turning off WiFi.
+                callback.waitForAvailable();
+            }
             SystemUtil.runShellCommand("svc wifi disable");
             if (wasWifiConnected) {
                 // Ensure we get both an onLost callback and a CONNECTIVITY_ACTION.
                 assertNotNull("Did not receive onLost callback after disabling wifi",
                         callback.waitForLost());
-            }
-            if (wasWifiConnected && expectLegacyBroadcast) {
-                assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                if (expectLegacyBroadcast) {
+                    assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                }
             }
         } catch (InterruptedException ex) {
             fail("disconnectFromWifi was interrupted");
@@ -599,12 +603,14 @@
 
         @Override
         public void onAvailable(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onAvailable " + network);
             currentNetwork = network;
             mAvailableCv.open();
         }
 
         @Override
         public void onLost(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onLost " + network);
             lastLostNetwork = network;
             if (network.equals(currentNetwork)) {
                 mAvailableCv.close();
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 448869d..bd30b77 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -93,6 +93,17 @@
     ]
 }
 
+// Subset of services-core used to by ConnectivityService tests to test VPN realistically.
+// This is stripped by jarjar (see rules below) from other unrelated classes, so tests do not
+// include most classes from services-core, which are unrelated and cause wrong code coverage
+// calculations.
+java_library {
+    name: "services.core-vpn",
+    static_libs: ["services.core"],
+    jarjar_rules: "vpn-jarjar-rules.txt",
+    visibility: ["//visibility:private"],
+}
+
 android_library {
     name: "FrameworksNetTestsLib",
     min_sdk_version: "30",
@@ -116,11 +127,11 @@
         "framework-protos",
         "mockito-target-minus-junit4",
         "net-tests-utils",
+        "net-utils-services-common",
         "platform-compat-test-rules",
         "platform-test-annotations",
         "service-connectivity-pre-jarjar",
-        "services.core",
-        "services.net",
+        "services.core-vpn",
     ],
     libs: [
         "android.net.ipsec.ike.stubs.module_lib",
@@ -144,6 +155,8 @@
     srcs: [":non-connectivity-module-test"],
     test_suites: ["device-tests"],
     static_libs: [
+        "services.core",
+        "services.net",
         "FrameworksNetTestsLib",
     ],
     libs: [
diff --git a/tests/unit/java/android/net/nsd/NsdManagerTest.java b/tests/unit/java/android/net/nsd/NsdManagerTest.java
index de77d23..30b8fcd 100644
--- a/tests/unit/java/android/net/nsd/NsdManagerTest.java
+++ b/tests/unit/java/android/net/nsd/NsdManagerTest.java
@@ -20,38 +20,32 @@
 import static libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.Context;
 import android.os.Build;
-import android.os.Handler;
-import android.os.HandlerThread;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.internal.util.AsyncChannel;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
-import com.android.testutils.HandlerUtils;
+import com.android.testutils.ExceptionUtils;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -67,9 +61,10 @@
 
     @Mock Context mContext;
     @Mock INsdManager mService;
-    MockServiceHandler mServiceHandler;
+    @Mock INsdServiceConnector mServiceConn;
 
     NsdManager mManager;
+    INsdManagerCallback mCallback;
 
     long mTimeoutMs = 200; // non-final so that tests can adjust the value.
 
@@ -77,91 +72,85 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
-        mServiceHandler = spy(MockServiceHandler.create(mContext));
-        doReturn(new Messenger(mServiceHandler)).when(mService).getMessenger();
-    }
-
-    @After
-    public void tearDown() throws Exception {
-        HandlerUtils.waitForIdle(mServiceHandler, mTimeoutMs);
-        mServiceHandler.chan.disconnect();
-        mServiceHandler.stop();
-        if (mManager != null) {
-            mManager.disconnect();
-        }
+        doReturn(mServiceConn).when(mService).connect(any());
+        mManager = new NsdManager(mContext, mService);
+        final ArgumentCaptor<INsdManagerCallback> cbCaptor = ArgumentCaptor.forClass(
+                INsdManagerCallback.class);
+        verify(mService).connect(cbCaptor.capture());
+        mCallback = cbCaptor.getValue();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testResolveServiceS() {
-        mManager = makeNsdManagerS();
+    public void testResolveServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestResolveService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testResolveServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testResolveServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestResolveService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testDiscoverServiceS() {
-        mManager = makeNsdManagerS();
+    public void testDiscoverServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestDiscoverService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testDiscoverServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testDiscoverServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestDiscoverService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testParallelResolveServiceS() {
-        mManager = makeNsdManagerS();
+    public void testParallelResolveServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestParallelResolveService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testParallelResolveServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testParallelResolveServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestParallelResolveService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testInvalidCallsS() {
-        mManager = makeNsdManagerS();
+    public void testInvalidCallsS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestInvalidCalls();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testInvalidCallsPreS() {
-        mManager = makeNsdManagerPreS();
+    public void testInvalidCallsPreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestInvalidCalls();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testRegisterServiceS() {
-        mManager = makeNsdManagerS();
+    public void testRegisterServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestRegisterService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testRegisterServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testRegisterServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestRegisterService();
     }
 
-    public void doTestResolveService() {
+    private void doTestResolveService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
@@ -169,18 +158,19 @@
         NsdManager.ResolveListener listener = mock(NsdManager.ResolveListener.class);
 
         manager.resolveService(request, listener);
-        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).resolveService(req.capture(), any()));
         int err = 33;
-        sendResponse(NsdManager.RESOLVE_SERVICE_FAILED, err, key1, null);
+        mCallback.onResolveServiceFailed(key1, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onResolveFailed(request, err);
 
         manager.resolveService(request, listener);
-        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).resolveService(req.capture(), any()));
+        mCallback.onResolveServiceSucceeded(key2, reply);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
     }
 
-    public void doTestParallelResolveService() {
+    private void doTestParallelResolveService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
@@ -190,19 +180,20 @@
         NsdManager.ResolveListener listener2 = mock(NsdManager.ResolveListener.class);
 
         manager.resolveService(request, listener1);
-        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).resolveService(req.capture(), any()));
 
         manager.resolveService(request, listener2);
-        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).resolveService(req.capture(), any()));
 
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key1, reply);
+        mCallback.onResolveServiceSucceeded(key2, reply);
+        mCallback.onResolveServiceSucceeded(key1, reply);
 
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
         verify(listener2, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
     }
 
-    public void doTestRegisterService() {
+    private void doTestRegisterService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request1 = new NsdServiceInfo("a_name", "a_type");
@@ -214,40 +205,43 @@
 
         // Register two services
         manager.registerService(request1, PROTOCOL, listener1);
-        int key1 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).registerService(req.capture(), any()));
 
         manager.registerService(request2, PROTOCOL, listener2);
-        int key2 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).registerService(req.capture(), any()));
 
         // First reques fails, second request succeeds
-        sendResponse(NsdManager.REGISTER_SERVICE_SUCCEEDED, 0, key2, request2);
+        mCallback.onRegisterServiceSucceeded(key2, request2);
         verify(listener2, timeout(mTimeoutMs).times(1)).onServiceRegistered(request2);
 
         int err = 1;
-        sendResponse(NsdManager.REGISTER_SERVICE_FAILED, err, key1, request1);
+        mCallback.onRegisterServiceFailed(key1, err);
         verify(listener1, timeout(mTimeoutMs).times(1)).onRegistrationFailed(request1, err);
 
         // Client retries first request, it succeeds
         manager.registerService(request1, PROTOCOL, listener1);
-        int key3 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key3 = getRequestKey(req ->
+                verify(mServiceConn, times(3)).registerService(req.capture(), any()));
 
-        sendResponse(NsdManager.REGISTER_SERVICE_SUCCEEDED, 0, key3, request1);
+        mCallback.onRegisterServiceSucceeded(key3, request1);
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceRegistered(request1);
 
         // First request is unregistered, it succeeds
         manager.unregisterService(listener1);
-        int key3again = verifyRequest(NsdManager.UNREGISTER_SERVICE);
+        int key3again = getRequestKey(req -> verify(mServiceConn).unregisterService(req.capture()));
         assertEquals(key3, key3again);
 
-        sendResponse(NsdManager.UNREGISTER_SERVICE_SUCCEEDED, 0, key3again, null);
+        mCallback.onUnregisterServiceSucceeded(key3again);
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceUnregistered(request1);
 
         // Second request is unregistered, it fails
         manager.unregisterService(listener2);
-        int key2again = verifyRequest(NsdManager.UNREGISTER_SERVICE);
+        int key2again = getRequestKey(req ->
+                verify(mServiceConn, times(2)).unregisterService(req.capture()));
         assertEquals(key2, key2again);
 
-        sendResponse(NsdManager.UNREGISTER_SERVICE_FAILED, err, key2again, null);
+        mCallback.onUnregisterServiceFailed(key2again, err);
         verify(listener2, timeout(mTimeoutMs).times(1)).onUnregistrationFailed(request2, err);
 
         // TODO: do not unregister listener until service is unregistered
@@ -260,7 +254,7 @@
         //verify(listener2, timeout(mTimeoutMs).times(1)).onServiceUnregistered(request2);
     }
 
-    public void doTestDiscoverService() {
+    private void doTestDiscoverService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo reply1 = new NsdServiceInfo("a_name", "a_type");
@@ -271,69 +265,73 @@
 
         // Client registers for discovery, request fails
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key1 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key1 = getRequestKey(req ->
+                verify(mServiceConn).discoverServices(req.capture(), any()));
 
         int err = 1;
-        sendResponse(NsdManager.DISCOVER_SERVICES_FAILED, err, key1, null);
+        mCallback.onDiscoverServicesFailed(key1, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onStartDiscoveryFailed("a_type", err);
 
         // Client retries, request succeeds
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key2 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).discoverServices(req.capture(), any()));
 
-        sendResponse(NsdManager.DISCOVER_SERVICES_STARTED, 0, key2, reply1);
+        mCallback.onDiscoverServicesStarted(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
 
         // mdns notifies about services
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply1);
+        mCallback.onServiceFound(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceFound(reply1);
 
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply2);
+        mCallback.onServiceFound(key2, reply2);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceFound(reply2);
 
-        sendResponse(NsdManager.SERVICE_LOST, 0, key2, reply2);
+        mCallback.onServiceLost(key2, reply2);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceLost(reply2);
 
 
         // Client unregisters its listener
         manager.stopServiceDiscovery(listener);
-        int key2again = verifyRequest(NsdManager.STOP_DISCOVERY);
+        int key2again = getRequestKey(req -> verify(mServiceConn).stopDiscovery(req.capture()));
         assertEquals(key2, key2again);
 
         // TODO: unregister listener immediately and stop notifying it about services
         // Notifications are still passed to the client's listener
-        sendResponse(NsdManager.SERVICE_LOST, 0, key2, reply1);
+        mCallback.onServiceLost(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceLost(reply1);
 
         // Client is notified of complete unregistration
-        sendResponse(NsdManager.STOP_DISCOVERY_SUCCEEDED, 0, key2again, "a_type");
+        mCallback.onStopDiscoverySucceeded(key2again);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStopped("a_type");
 
         // Notifications are not passed to the client anymore
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply3);
+        mCallback.onServiceFound(key2, reply3);
         verify(listener, timeout(mTimeoutMs).times(0)).onServiceLost(reply3);
 
 
         // Client registers for service discovery
         reset(listener);
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key3 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key3 = getRequestKey(req ->
+                verify(mServiceConn, times(3)).discoverServices(req.capture(), any()));
 
-        sendResponse(NsdManager.DISCOVER_SERVICES_STARTED, 0, key3, reply1);
+        mCallback.onDiscoverServicesStarted(key3, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
         // Client unregisters immediately, it fails
         manager.stopServiceDiscovery(listener);
-        int key3again = verifyRequest(NsdManager.STOP_DISCOVERY);
+        int key3again = getRequestKey(req ->
+                verify(mServiceConn, times(2)).stopDiscovery(req.capture()));
         assertEquals(key3, key3again);
 
         err = 2;
-        sendResponse(NsdManager.STOP_DISCOVERY_FAILED, err, key3again, "a_type");
+        mCallback.onStopDiscoveryFailed(key3again, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onStopDiscoveryFailed("a_type", err);
 
         // New notifications are not passed to the client anymore
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key3, reply1);
+        mCallback.onServiceFound(key3, reply1);
         verify(listener, timeout(mTimeoutMs).times(0)).onServiceFound(reply1);
     }
 
@@ -398,77 +396,10 @@
         }
     }
 
-    NsdManager makeNsdManagerS() {
-        // Expect we'll get 2 AsyncChannel related msgs.
-        return makeManager(2);
-    }
-
-    NsdManager makeNsdManagerPreS() {
-        // Expect we'll get 3 msgs. 2 AsyncChannel related msgs + 1 additional daemon startup msg.
-        return makeManager(3);
-    }
-
-    NsdManager makeManager(int expectedMsgCount) {
-        NsdManager manager = new NsdManager(mContext, mService);
-        // Acknowledge first two messages connecting the AsyncChannel.
-        verify(mServiceHandler, timeout(mTimeoutMs).times(expectedMsgCount)).handleMessage(any());
-
-        reset(mServiceHandler);
-        assertNotNull(mServiceHandler.chan);
-        return manager;
-    }
-
-    int verifyRequest(int expectedMessageType) {
-        HandlerUtils.waitForIdle(mServiceHandler, mTimeoutMs);
-        verify(mServiceHandler, timeout(mTimeoutMs)).handleMessage(any());
-        reset(mServiceHandler);
-        Message received = mServiceHandler.getLastMessage();
-        assertEquals(NsdManager.nameOf(expectedMessageType), NsdManager.nameOf(received.what));
-        return received.arg2;
-    }
-
-    void sendResponse(int replyType, int arg, int key, Object obj) {
-        mServiceHandler.chan.sendMessage(replyType, arg, key, obj);
-    }
-
-    // Implements the server side of AsyncChannel connection protocol
-    public static class MockServiceHandler extends Handler {
-        public final Context context;
-        public AsyncChannel chan;
-        public Message lastMessage;
-
-        MockServiceHandler(Looper l, Context c) {
-            super(l);
-            context = c;
-        }
-
-        synchronized Message getLastMessage() {
-            return lastMessage;
-        }
-
-        synchronized void setLastMessage(Message msg) {
-            lastMessage = obtainMessage();
-            lastMessage.copyFrom(msg);
-        }
-
-        @Override
-        public void handleMessage(Message msg) {
-            setLastMessage(msg);
-            if (msg.what == AsyncChannel.CMD_CHANNEL_FULL_CONNECTION) {
-                chan = new AsyncChannel();
-                chan.connect(context, this, msg.replyTo);
-                chan.sendMessage(AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED);
-            }
-        }
-
-        void stop() {
-            getLooper().quitSafely();
-        }
-
-        static MockServiceHandler create(Context context) {
-            HandlerThread t = new HandlerThread("mock-service-handler");
-            t.start();
-            return new MockServiceHandler(t.getLooper(), context);
-        }
+    int getRequestKey(ExceptionUtils.ThrowingConsumer<ArgumentCaptor<Integer>> verifier)
+            throws Exception {
+        final ArgumentCaptor<Integer> captor = ArgumentCaptor.forClass(Integer.class);
+        verifier.accept(captor);
+        return captor.getValue();
     }
 }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index e0dd886..e76762e 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -249,6 +249,8 @@
 import android.net.NetworkStateSnapshot;
 import android.net.NetworkTestResultParcelable;
 import android.net.OemNetworkPreferences;
+import android.net.PacProxyManager;
+import android.net.Proxy;
 import android.net.ProxyInfo;
 import android.net.QosCallbackException;
 import android.net.QosFilter;
@@ -478,6 +480,7 @@
     private Context mContext;
     private NetworkPolicyCallback mPolicyCallback;
     private WrappedMultinetworkPolicyTracker mPolicyTracker;
+    private ProxyTracker mProxyTracker;
     private HandlerThread mAlarmManagerThread;
     private TestNetIdManager mNetIdManager;
     private QosCallbackMockHelper mQosCallbackMockHelper;
@@ -511,7 +514,7 @@
     @Mock VpnProfileStore mVpnProfileStore;
     @Mock SystemConfigManager mSystemConfigManager;
     @Mock Resources mResources;
-    @Mock ProxyTracker mProxyTracker;
+    @Mock PacProxyManager mPacProxyManager;
 
     // BatteryStatsManager is final and cannot be mocked with regular mockito, so just mock the
     // underlying binder calls.
@@ -613,6 +616,7 @@
             if (Context.SYSTEM_CONFIG_SERVICE.equals(name)) return mSystemConfigManager;
             if (Context.NETWORK_STATS_SERVICE.equals(name)) return mStatsManager;
             if (Context.BATTERY_STATS_SERVICE.equals(name)) return mBatteryStatsManager;
+            if (Context.PAC_PROXY_SERVICE.equals(name)) return mPacProxyManager;
             return super.getSystemService(name);
         }
 
@@ -1703,6 +1707,8 @@
 
         mCsHandlerThread = new HandlerThread("TestConnectivityService");
         mVMSHandlerThread = new HandlerThread("TestVpnManagerService");
+        mProxyTracker = new ProxyTracker(mServiceContext, mock(Handler.class),
+                16 /* EVENT_PROXY_HAS_CHANGED */);
 
         initMockedResources();
         final Context mockResContext = mock(Context.class);
@@ -2135,6 +2141,39 @@
         return expected;
     }
 
+    private ExpectedBroadcast expectProxyChangeAction(ProxyInfo proxy) {
+        return registerPacProxyBroadcastThat(intent -> {
+            final ProxyInfo actualProxy = (ProxyInfo) intent.getExtra(Proxy.EXTRA_PROXY_INFO,
+                    ProxyInfo.buildPacProxy(Uri.EMPTY));
+            return proxy.equals(actualProxy);
+        });
+    }
+
+    private ExpectedBroadcast registerPacProxyBroadcast() {
+        return registerPacProxyBroadcastThat(intent -> true);
+    }
+
+    private ExpectedBroadcast registerPacProxyBroadcastThat(
+            @NonNull final Predicate<Intent> filter) {
+        final IntentFilter intentFilter = new IntentFilter(Proxy.PROXY_CHANGE_ACTION);
+        // AtomicReference allows receiver to access expected even though it is constructed later.
+        final AtomicReference<ExpectedBroadcast> expectedRef = new AtomicReference<>();
+        final BroadcastReceiver receiver = new BroadcastReceiver() {
+            public void onReceive(Context context, Intent intent) {
+                final ProxyInfo proxy = (ProxyInfo) intent.getExtra(
+                            Proxy.EXTRA_PROXY_INFO, ProxyInfo.buildPacProxy(Uri.EMPTY));
+                Log.d(TAG, "Receive PROXY_CHANGE_ACTION, proxy = " + proxy);
+                if (filter.test(intent)) {
+                    expectedRef.get().complete(intent);
+                }
+            }
+        };
+        final ExpectedBroadcast expected = new ExpectedBroadcast(receiver);
+        expectedRef.set(expected);
+        mServiceContext.registerReceiver(receiver, intentFilter);
+        return expected;
+    }
+
     private boolean extraInfoInBroadcastHasExpectedNullness(NetworkInfo ni) {
         final DetailedState state = ni.getDetailedState();
         if (state == DetailedState.CONNECTED && ni.getExtraInfo() == null) return false;
@@ -9945,7 +9984,7 @@
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         final Network wifiNetwork = mWiFiNetworkAgent.getNetwork();
-        doReturn(testProxyInfo).when(mService.mProxyTracker).getGlobalProxy();
+        mProxyTracker.setGlobalProxy(testProxyInfo);
         assertEquals(testProxyInfo, mService.getProxyForNetwork(wifiNetwork));
     }
 
@@ -11353,6 +11392,7 @@
         assertNull(mService.getProxyForNetwork(null));
         assertNull(mCm.getDefaultProxy());
 
+        final ExpectedBroadcast b1 = registerPacProxyBroadcast();
         final LinkProperties lp = new LinkProperties();
         lp.setInterfaceName("tun0");
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
@@ -11362,9 +11402,10 @@
         mMockVpn.establish(lp, VPN_UID, vpnRanges);
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);
         // VPN is connected but proxy is not set, so there is no need to send proxy broadcast.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b1.expectNoBroadcast(500);
 
         // Update to new range which is old range minus APP1, i.e. only APP2
+        final ExpectedBroadcast b2 = registerPacProxyBroadcast();
         final Set<UidRange> newRanges = new HashSet<>(asList(
                 new UidRange(vpnRange.start, APP1_UID - 1),
                 new UidRange(APP1_UID + 1, vpnRange.stop)));
@@ -11375,37 +11416,37 @@
         assertVpnUidRangesUpdated(false, vpnRanges, VPN_UID);
 
         // Uid has changed but proxy is not set, so there is no need to send proxy broadcast.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b2.expectNoBroadcast(500);
 
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
+        final ExpectedBroadcast b3 = registerPacProxyBroadcast();
         lp.setHttpProxy(testProxyInfo);
         mMockVpn.sendLinkProperties(lp);
         waitForIdle();
         // Proxy is set, so send a proxy broadcast.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b3.expectBroadcast();
 
+        final ExpectedBroadcast b4 = registerPacProxyBroadcast();
         mMockVpn.setUids(vpnRanges);
         waitForIdle();
         // Uid has changed and proxy is already set, so send a proxy broadcast.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b4.expectBroadcast();
 
+        final ExpectedBroadcast b5 = registerPacProxyBroadcast();
         // Proxy is removed, send a proxy broadcast.
         lp.setHttpProxy(null);
         mMockVpn.sendLinkProperties(lp);
         waitForIdle();
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b5.expectBroadcast();
 
         // Proxy is added in WiFi(default network), setDefaultProxy will be called.
         final LinkProperties wifiLp = mCm.getLinkProperties(mWiFiNetworkAgent.getNetwork());
         assertNotNull(wifiLp);
+        final ExpectedBroadcast b6 = expectProxyChangeAction(testProxyInfo);
         wifiLp.setHttpProxy(testProxyInfo);
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         waitForIdle();
-        verify(mProxyTracker, times(1)).setDefaultProxy(eq(testProxyInfo));
-        reset(mProxyTracker);
+        b6.expectBroadcast();
     }
 
     @Test
@@ -11424,18 +11465,21 @@
         lp.setHttpProxy(testProxyInfo);
         final UidRange vpnRange = PRIMARY_UIDRANGE;
         final Set<UidRange> vpnRanges = Collections.singleton(vpnRange);
+        final ExpectedBroadcast b1 = registerPacProxyBroadcast();
         mMockVpn.setOwnerAndAdminUid(VPN_UID);
         mMockVpn.registerAgent(false, vpnRanges, lp);
         // In any case, the proxy broadcast won't be sent before VPN goes into CONNECTED state.
         // Otherwise, the app that calls ConnectivityManager#getDefaultProxy() when it receives the
         // proxy broadcast will get null.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b1.expectNoBroadcast(500);
+
+        final ExpectedBroadcast b2 = registerPacProxyBroadcast();
         mMockVpn.connect(true /* validated */, true /* hasInternet */, false /* isStrictMode */);
         waitForIdle();
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);
         // Vpn is connected with proxy, so the proxy broadcast will be sent to inform the apps to
         // update their proxy data.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
+        b2.expectBroadcast();
     }
 
     @Test
@@ -11464,10 +11508,10 @@
         final LinkProperties cellularLp = new LinkProperties();
         cellularLp.setInterfaceName(MOBILE_IFNAME);
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
+        final ExpectedBroadcast b = registerPacProxyBroadcast();
         cellularLp.setHttpProxy(testProxyInfo);
         mCellNetworkAgent.sendLinkProperties(cellularLp);
-        waitForIdle();
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
+        b.expectBroadcast();
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 4d2970a..4172553 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -20,6 +20,7 @@
 import static libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -33,14 +34,19 @@
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.ContentResolver;
 import android.content.Context;
+import android.net.nsd.INsdManagerCallback;
+import android.net.nsd.INsdServiceConnector;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
+import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
 import com.android.server.NsdService.DaemonConnection;
@@ -56,11 +62,15 @@
 import org.junit.Test;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
+import org.mockito.AdditionalAnswers;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.Spy;
 
+import java.util.LinkedList;
+import java.util.Queue;
+
 // TODOs:
 //  - test client can send requests and receive replies
 //  - test NSD_ON ENABLE/DISABLED listening
@@ -73,6 +83,11 @@
     private static final long CLEANUP_DELAY_MS = 500;
     private static final long TIMEOUT_MS = 500;
 
+    // Records INsdManagerCallback created when NsdService#connect is called.
+    // Only accessed on the test thread, since NsdService#connect is called by the NsdManager
+    // constructor called on the test thread.
+    private final Queue<INsdManagerCallback> mCreatedCallbacks = new LinkedList<>();
+
     @Rule
     public TestRule compatChangeRule = new PlatformCompatChangeRule();
     @Mock Context mContext;
@@ -83,6 +98,16 @@
     HandlerThread mThread;
     TestHandler mHandler;
 
+    private static class LinkToDeathRecorder extends Binder {
+        IBinder.DeathRecipient mDr;
+
+        @Override
+        public void linkToDeath(@NonNull DeathRecipient recipient, int flags) {
+            super.linkToDeath(recipient, flags);
+            mDr = recipient;
+        }
+    }
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -103,26 +128,30 @@
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testPreSClients() {
+    public void testPreSClients() throws Exception {
         when(mSettings.isEnabled()).thenReturn(true);
         NsdService service = makeService();
 
         // Pre S client connected, the daemon should be started.
-        NsdManager client1 = connectClient(service);
+        connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient1 = verifyLinkToDeath(cb1);
         verify(mDaemon, times(1)).maybeStart();
         verifyDaemonCommands("start-service");
 
-        NsdManager client2 = connectClient(service);
+        connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb2 = getCallback();
+        final IBinder.DeathRecipient deathRecipient2 = verifyLinkToDeath(cb2);
         verify(mDaemon, times(1)).maybeStart();
 
-        client1.disconnect();
+        deathRecipient1.binderDied();
         // Still 1 client remains, daemon shouldn't be stopped.
         waitForIdle();
         verify(mDaemon, never()).maybeStop();
 
-        client2.disconnect();
+        deathRecipient2.binderDied();
         // All clients are disconnected, the daemon should be stopped.
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         verifyDaemonCommands("stop-service");
@@ -130,43 +159,51 @@
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testNoDaemonStartedWhenClientsConnect() {
+    public void testNoDaemonStartedWhenClientsConnect() throws Exception {
         when(mSettings.isEnabled()).thenReturn(true);
-
-        NsdService service = makeService();
+        final NsdService service = makeService();
 
         // Creating an NsdManager will not cause any cmds executed, which means
         // no daemon is started.
-        NsdManager client1 = connectClient(service);
+        connectClient(service);
         waitForIdle();
         verify(mDaemon, never()).execute(any());
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient1 = verifyLinkToDeath(cb1);
 
         // Creating another NsdManager will not cause any cmds executed.
-        NsdManager client2 = connectClient(service);
+        connectClient(service);
         waitForIdle();
         verify(mDaemon, never()).execute(any());
+        final INsdManagerCallback cb2 = getCallback();
+        final IBinder.DeathRecipient deathRecipient2 = verifyLinkToDeath(cb2);
 
         // If there is no active request, try to clean up the daemon
         // every time the client disconnects.
-        client1.disconnect();
+        deathRecipient1.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         reset(mDaemon);
-        client2.disconnect();
+        deathRecipient2.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
+    }
 
-        client1.disconnect();
-        client2.disconnect();
+    private IBinder.DeathRecipient verifyLinkToDeath(INsdManagerCallback cb)
+            throws Exception {
+        final IBinder.DeathRecipient dr = ((LinkToDeathRecorder) cb.asBinder()).mDr;
+        assertNotNull(dr);
+        return dr;
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testClientRequestsAreGCedAtDisconnection() {
+    public void testClientRequestsAreGCedAtDisconnection() throws Exception {
         when(mSettings.isEnabled()).thenReturn(true);
-
         NsdService service = makeService();
-        NsdManager client = connectClient(service);
 
+        NsdManager client = connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient = verifyLinkToDeath(cb1);
         verify(mDaemon, never()).maybeStart();
         verify(mDaemon, never()).execute(any());
 
@@ -195,18 +232,16 @@
         verifyDaemonCommand("resolve 4 a_name a_type local.");
 
         // Client disconnects, stop the daemon after CLEANUP_DELAY_MS.
-        client.disconnect();
+        deathRecipient.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         // checks that request are cleaned
         verifyDaemonCommands("stop-register 2", "stop-discover 3",
                 "stop-resolve 4", "stop-service");
-
-        client.disconnect();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testCleanupDelayNoRequestActive() {
+    public void testCleanupDelayNoRequestActive() throws Exception {
         when(mSettings.isEnabled()).thenReturn(true);
 
         NsdService service = makeService();
@@ -218,6 +253,8 @@
         client.registerService(request, PROTOCOL, listener1);
         waitForIdle();
         verify(mDaemon, times(1)).maybeStart();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient = verifyLinkToDeath(cb1);
         verifyDaemonCommands("start-service", "register 2 a_name a_type 2201");
 
         client.unregisterService(listener1);
@@ -226,7 +263,7 @@
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         verifyDaemonCommand("stop-service");
         reset(mDaemon);
-        client.disconnect();
+        deathRecipient.binderDied();
         // Client disconnects, after CLEANUP_DELAY_MS, maybeStop the daemon.
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
     }
@@ -240,12 +277,29 @@
             mDaemonCallback = callback;
             return mDaemon;
         };
-        NsdService service = new NsdService(mContext, mSettings,
-                mHandler, supplier, CLEANUP_DELAY_MS);
+        final NsdService service = new NsdService(mContext, mSettings,
+                mHandler, supplier, CLEANUP_DELAY_MS) {
+            @Override
+            public INsdServiceConnector connect(INsdManagerCallback baseCb) {
+                // Wrap the callback in a transparent mock, to mock asBinder returning a
+                // LinkToDeathRecorder. This will allow recording the binder death recipient
+                // registered on the callback. Use a transparent mock and not a spy as the actual
+                // implementation class is not public and cannot be spied on by Mockito.
+                final INsdManagerCallback cb = mock(INsdManagerCallback.class,
+                        AdditionalAnswers.delegatesTo(baseCb));
+                doReturn(new LinkToDeathRecorder()).when(cb).asBinder();
+                mCreatedCallbacks.add(cb);
+                return super.connect(cb);
+            }
+        };
         verify(mDaemon, never()).execute(any(String.class));
         return service;
     }
 
+    private INsdManagerCallback getCallback() {
+        return mCreatedCallbacks.remove();
+    }
+
     NsdManager connectClient(NsdService service) {
         return new NsdManager(mContext, service);
     }
diff --git a/tests/unit/vpn-jarjar-rules.txt b/tests/unit/vpn-jarjar-rules.txt
new file mode 100644
index 0000000..16661b9
--- /dev/null
+++ b/tests/unit/vpn-jarjar-rules.txt
@@ -0,0 +1,4 @@
+# Only keep classes imported by ConnectivityServiceTest
+keep com.android.server.VpnManagerService
+keep com.android.server.connectivity.Vpn
+keep com.android.server.connectivity.VpnProfileStore
\ No newline at end of file