Fix flakes due to tests resetting CarrierConfig

Some tests use CarrierConfigManager#overrideConfig(subId, bundle) to set
some configs, and then overrideConfig(subId, null) to reset the
override.  However this also resets overrides that may have been set by
the instrumentation, such as ConnectivityTestTargetPreparer /
CarrierConfigSetupTest.

Avoid this by providing a common utility to set and tear down carrier
config, which reset the overrides to the original setting value instead
of clearing all overrides.

Test: atest
Bug: 253698734
Change-Id: Ia25b28d551b62205c52f91e0ffad1b94710c0828
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/CarrierConfigRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/CarrierConfigRule.kt
new file mode 100644
index 0000000..a93ae3e
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/CarrierConfigRule.kt
@@ -0,0 +1,127 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils.com.android.testutils
+
+import android.Manifest.permission.MODIFY_PHONE_STATE
+import android.Manifest.permission.READ_PHONE_STATE
+import android.os.PersistableBundle
+import android.telephony.CarrierConfigManager
+import android.util.Log
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.modules.utils.build.SdkLevel.isAtLeastU
+import com.android.testutils.runAsShell
+import com.android.testutils.tryTest
+import kotlin.test.assertNotNull
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+private val TAG = CarrierConfigRule::class.simpleName
+
+/**
+ * A [TestRule] that helps set [CarrierConfigManager] overrides for tests and clean up the test
+ * configuration automatically on teardown.
+ */
+class CarrierConfigRule : TestRule {
+    private val ccm by lazy { InstrumentationRegistry.getInstrumentation().context.getSystemService(
+        CarrierConfigManager::class.java
+    ) }
+
+    // Map of (subId) -> (original values of overridden settings)
+    private val originalConfigs = mutableMapOf<Int, PersistableBundle>()
+
+    override fun apply(base: Statement, description: Description): Statement {
+        return CarrierConfigStatement(base, description)
+    }
+
+    private inner class CarrierConfigStatement(
+        private val base: Statement,
+        private val description: Description
+    ) : Statement() {
+        override fun evaluate() {
+            tryTest {
+                base.evaluate()
+            } cleanup {
+                cleanUpNow()
+            }
+        }
+    }
+
+    /**
+     * Add carrier config overrides with the specified configuration.
+     *
+     * The overrides will automatically be cleaned up when the test case finishes.
+     */
+    fun addConfigOverrides(subId: Int, config: PersistableBundle) {
+        val originalConfig = originalConfigs.computeIfAbsent(subId) { PersistableBundle() }
+        val overrideKeys = config.keySet()
+        val previousValues = runAsShell(READ_PHONE_STATE) {
+            ccm.getConfigForSubIdCompat(subId, overrideKeys)
+        }
+        // If a key is already in the originalConfig, keep the oldest original overrides
+        originalConfig.keySet().forEach {
+            previousValues.remove(it)
+        }
+        originalConfig.putAll(previousValues)
+
+        runAsShell(MODIFY_PHONE_STATE) {
+            ccm.overrideConfig(subId, config)
+        }
+    }
+
+    /**
+     * Cleanup overrides that were added by the test case.
+     *
+     * This will be called automatically on test teardown, so it does not need to be called by the
+     * test case unless cleaning up earlier is required.
+     */
+    fun cleanUpNow() {
+        runAsShell(MODIFY_PHONE_STATE) {
+            originalConfigs.forEach { (subId, config) ->
+                try {
+                    // Do not use overrideConfig with null, as it would reset configs that may
+                    // have been set by target preparers such as
+                    // ConnectivityTestTargetPreparer / CarrierConfigSetupTest.
+                    ccm.overrideConfig(subId, config)
+                } catch (e: Throwable) {
+                    Log.e(TAG, "Error resetting carrier config for subId $subId")
+                }
+            }
+            originalConfigs.clear()
+        }
+    }
+}
+
+private fun CarrierConfigManager.getConfigForSubIdCompat(
+    subId: Int,
+    keys: Set<String>
+): PersistableBundle {
+    return if (isAtLeastU()) {
+        // This method is U+
+        getConfigForSubId(subId, *keys.toTypedArray())
+    } else {
+        @Suppress("DEPRECATION")
+        val config = assertNotNull(getConfigForSubId(subId))
+        val allKeys = config.keySet().toList()
+        allKeys.forEach {
+            if (!keys.contains(it)) {
+                config.remove(it)
+            }
+        }
+        config
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
index ceb48d4..faaadee 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityDiagnosticsManagerTest.java
@@ -88,9 +88,11 @@
 import com.android.net.module.util.ArrayTrackRecord;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.DevSdkIgnoreRunner;
+import com.android.testutils.com.android.testutils.CarrierConfigRule;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -110,6 +112,9 @@
 public class ConnectivityDiagnosticsManagerTest {
     private static final String TAG = ConnectivityDiagnosticsManagerTest.class.getSimpleName();
 
+    @Rule
+    public final CarrierConfigRule mCarrierConfigRule = new CarrierConfigRule();
+
     private static final int CALLBACK_TIMEOUT_MILLIS = 5000;
     private static final int NO_CALLBACK_INVOKED_TIMEOUT = 500;
     private static final long TIMESTAMP = 123456789L;
@@ -264,9 +269,6 @@
             doBroadcastCarrierConfigsAndVerifyOnConnectivityReportAvailable(
                     subId, carrierConfigReceiver, testNetworkCallback);
         }, () -> {
-                runWithShellPermissionIdentity(
-                    () -> mCarrierConfigManager.overrideConfig(subId, null),
-                    android.Manifest.permission.MODIFY_PHONE_STATE);
             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback);
             mContext.unregisterReceiver(carrierConfigReceiver);
             });
@@ -291,9 +293,9 @@
                 CarrierConfigManager.KEY_CARRIER_CERTIFICATE_STRING_ARRAY,
                 new String[] {getCertHashForThisPackage()});
 
+        mCarrierConfigRule.addConfigOverrides(subId, carrierConfigs);
         runWithShellPermissionIdentity(
                 () -> {
-                    mCarrierConfigManager.overrideConfig(subId, carrierConfigs);
                     mCarrierConfigManager.notifyConfigChangedForSubId(subId);
                 },
                 android.Manifest.permission.MODIFY_PHONE_STATE);
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index ae848d5..286e08c 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -140,6 +140,7 @@
 import com.android.testutils.TestableNetworkAgent.CallbackEntry.OnValidationStatus
 import com.android.testutils.TestableNetworkCallback
 import com.android.testutils.assertThrows
+import com.android.testutils.com.android.testutils.CarrierConfigRule
 import com.android.testutils.runAsShell
 import com.android.testutils.tryTest
 import com.android.testutils.waitForIdle
@@ -167,6 +168,7 @@
 import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.any
@@ -225,6 +227,9 @@
 @IgnoreUpTo(Build.VERSION_CODES.R)
 @RunWith(DevSdkIgnoreRunner::class)
 class NetworkAgentTest {
+    @get:Rule
+    val carrierConfigRule = CarrierConfigRule()
+
     private val LOCAL_IPV4_ADDRESS = InetAddresses.parseNumericAddress("192.0.2.1")
     private val REMOTE_IPV4_ADDRESS = InetAddresses.parseNumericAddress("192.0.2.2")
 
@@ -714,7 +719,6 @@
         }
 
         val tm = realContext.getSystemService(TelephonyManager::class.java)!!
-        val ccm = realContext.getSystemService(CarrierConfigManager::class.java)!!
 
         val cv = ConditionVariable()
         val cpb = PrivilegeWaiterCallback(cv)
@@ -734,16 +738,13 @@
                 return@tryTest
             }
             cv.close()
-            runAsShell(MODIFY_PHONE_STATE) {
-                val carrierConfigs = if (hold) {
-                    PersistableBundle().also {
-                        it.putStringArray(CarrierConfigManager.KEY_CARRIER_CERTIFICATE_STRING_ARRAY,
-                                arrayOf(getCertHash()))
-                    }
-                } else {
-                    null
-                }
-                ccm.overrideConfig(subId, carrierConfigs)
+            if (hold) {
+                carrierConfigRule.addConfigOverrides(subId, PersistableBundle().also {
+                    it.putStringArray(CarrierConfigManager.KEY_CARRIER_CERTIFICATE_STRING_ARRAY,
+                        arrayOf(getCertHash()))
+                })
+            } else {
+                carrierConfigRule.cleanUpNow()
             }
             assertTrue(cv.block(DEFAULT_TIMEOUT_MS), "Can't change carrier privilege")
         } cleanup {
diff --git a/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java b/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
index e0e22e5..6d53ddf 100644
--- a/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
+++ b/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
@@ -84,9 +84,11 @@
 
 import com.android.modules.utils.build.SdkLevel;
 import com.android.testutils.ParcelUtils;
+import com.android.testutils.com.android.testutils.CarrierConfigRule;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -102,6 +104,8 @@
 
 @RunWith(AndroidJUnit4.class)
 public class TetheringManagerTest {
+    @Rule
+    public final CarrierConfigRule mCarrierConfigRule = new CarrierConfigRule();
 
     private Context mContext;
 
@@ -527,22 +531,13 @@
         // Override carrier config to ignore entitlement check.
         final PersistableBundle bundle = new PersistableBundle();
         bundle.putBoolean(CarrierConfigManager.KEY_REQUIRE_ENTITLEMENT_CHECKS_BOOL, false);
-        overrideCarrierConfig(bundle);
+        mCarrierConfigRule.addConfigOverrides(
+                SubscriptionManager.getDefaultSubscriptionId(), bundle);
 
         // Verify that requestLatestTetheringEntitlementResult() can get entitlement
         // result TETHER_ERROR_NO_ERROR due to provisioning bypassed.
         assertEntitlementResult(listener -> mTM.requestLatestTetheringEntitlementResult(
                 TETHERING_WIFI, false, c -> c.run(), listener), TETHER_ERROR_NO_ERROR);
-
-        // Reset carrier config.
-        overrideCarrierConfig(null);
-    }
-
-    private void overrideCarrierConfig(PersistableBundle bundle) {
-        final CarrierConfigManager configManager = (CarrierConfigManager) mContext
-                .getSystemService(Context.CARRIER_CONFIG_SERVICE);
-        final int subId = SubscriptionManager.getDefaultSubscriptionId();
-        runAsShell(MODIFY_PHONE_STATE, () -> configManager.overrideConfig(subId, bundle));
     }
 
     private boolean isTetheringApnRequired() {