diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
index 5ae9232..a5a1aeb 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
@@ -131,4 +131,8 @@
         }
         mgr.unregisterNetworkCallback(cb)
     }
+
+    fun disconnect() {
+        agent.unregister()
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index 68613a6..9332435 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -24,13 +24,16 @@
 import android.net.PacProxyManager
 import android.net.RouteInfo
 import android.net.networkstack.NetworkStackClientBase
+import android.os.BatteryStatsManager
 import android.os.Handler
 import android.os.HandlerThread
 import android.os.UserHandle
 import android.os.UserManager
 import android.telephony.TelephonyManager
 import android.testing.TestableContext
+import android.util.ArraySet
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.internal.app.IBatteryStats
 import com.android.internal.util.test.BroadcastInterceptingContext
 import com.android.modules.utils.build.SdkLevel
 import com.android.networkstack.apishim.common.UnsupportedApiLevelException
@@ -41,6 +44,7 @@
 import com.android.server.connectivity.MultinetworkPolicyTracker
 import com.android.server.connectivity.MultinetworkPolicyTrackerTestDependencies
 import com.android.server.connectivity.ProxyTracker
+import com.android.testutils.visibleOnHandlerThread
 import com.android.testutils.waitForIdle
 import org.mockito.AdditionalAnswers.delegatesTo
 import org.mockito.Mockito.doAnswer
@@ -71,7 +75,7 @@
     init {
         if (!SdkLevel.isAtLeastS()) {
             throw UnsupportedApiLevelException("CSTest subclasses must be annotated to only " +
-                    "run on S+, e.g. @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.R)");
+                    "run on S+, e.g. @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.R)")
         }
     }
 
@@ -112,6 +116,7 @@
     val proxyTracker = ProxyTracker(context, mock<Handler>(), 16 /* EVENT_PROXY_HAS_CHANGED */)
     val alarmManager = makeMockAlarmManager()
     val systemConfigManager = makeMockSystemConfigManager()
+    val batteryManager = BatteryStatsManager(mock<IBatteryStats>())
     val telephonyManager = mock<TelephonyManager>().also {
         doReturn(true).`when`(it).isDataCapable()
     }
@@ -150,6 +155,25 @@
         // checking permissions.
         override fun isFeatureEnabled(context: Context?, name: String?) =
                 enabledFeatures[name] ?: fail("Unmocked feature $name, see CSTest.enabledFeatures")
+
+        // Mocked change IDs
+        private val enabledChangeIds = ArraySet<Long>()
+        fun setChangeIdEnabled(enabled: Boolean, changeId: Long) {
+            // enabledChangeIds is read on the handler thread and maybe the test thread, so
+            // make sure both threads see it before continuing.
+            visibleOnHandlerThread(csHandler) {
+                if (enabled) {
+                    enabledChangeIds.add(changeId)
+                } else {
+                    enabledChangeIds.remove(changeId)
+                }
+            }
+        }
+
+        override fun isChangeEnabled(changeId: Long, pkg: String, user: UserHandle) =
+                changeId in enabledChangeIds
+        override fun isChangeEnabled(changeId: Long, uid: Int) =
+                changeId in enabledChangeIds
     }
 
     inner class CSContext(base: Context) : BroadcastInterceptingContext(base) {
@@ -196,6 +220,7 @@
             Context.ACTIVITY_SERVICE -> activityManager
             Context.SYSTEM_CONFIG_SERVICE -> systemConfigManager
             Context.TELEPHONY_SERVICE -> telephonyManager
+            Context.BATTERY_STATS_SERVICE -> batteryManager
             Context.STATS_MANAGER -> null // Stats manager is final and can't be mocked
             else -> super.getSystemService(serviceName)
         }