Fix flaky TeleServices tests causing module errors

Some tests were using reflection to replace static instances but not
resetting the previous values, causing the device to remain in a bad
state. Add a replaceInstance and restoreInstances method to
TelephonyTestBase and have these classes extend it so cleanup can be
done in the @After method.
Fix flaky NotificationMgrTests causing module errors due to incomplete
mocks.
Also revert changes added to help debug b/256244889.

Test: atest TeleServiceTests
Bug: 256244889
Change-Id: Id758e207da862f450d5c90971cb3ea1bd2afe40d
diff --git a/tests/src/com/android/TelephonyTestBase.java b/tests/src/com/android/TelephonyTestBase.java
index ffda81b..d72d85e 100644
--- a/tests/src/com/android/TelephonyTestBase.java
+++ b/tests/src/com/android/TelephonyTestBase.java
@@ -24,10 +24,16 @@
 
 import com.android.internal.telephony.PhoneConfigurationManager;
 
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
@@ -40,6 +46,10 @@
 
     protected TestContext mContext;
 
+    private final HashMap<InstanceKey, Object> mOldInstances = new HashMap<>();
+    private final LinkedList<InstanceKey> mInstanceKeys = new LinkedList<>();
+
+    @Before
     public void setUp() throws Exception {
         mContext = spy(new TestContext());
         // Set up the looper if it does not exist on the test thread.
@@ -56,9 +66,11 @@
         }
     }
 
+    @After
     public void tearDown() throws Exception {
         // Ensure there are no static references to handlers after test completes.
         PhoneConfigurationManager.unregisterAllMultiSimConfigChangeRegistrants();
+        restoreInstances();
     }
 
     protected final boolean waitForExecutorAction(Executor executor, long timeoutMillis) {
@@ -108,6 +120,61 @@
         }
     }
 
+    protected synchronized void replaceInstance(final Class c, final String instanceName,
+            final Object obj, final Object newValue)
+            throws Exception {
+        Field field = c.getDeclaredField(instanceName);
+        field.setAccessible(true);
+
+        InstanceKey key = new InstanceKey(c, instanceName, obj);
+        if (!mOldInstances.containsKey(key)) {
+            mOldInstances.put(key, field.get(obj));
+            mInstanceKeys.add(key);
+        }
+        field.set(obj, newValue);
+    }
+
+    private synchronized void restoreInstances() throws Exception {
+        Iterator<InstanceKey> it = mInstanceKeys.descendingIterator();
+
+        while (it.hasNext()) {
+            InstanceKey key = it.next();
+            Field field = key.mClass.getDeclaredField(key.mInstName);
+            field.setAccessible(true);
+            field.set(key.mObj, mOldInstances.get(key));
+        }
+
+        mInstanceKeys.clear();
+        mOldInstances.clear();
+    }
+
+    private static class InstanceKey {
+        public final Class mClass;
+        public final String mInstName;
+        public final Object mObj;
+        InstanceKey(final Class c, final String instName, final Object obj) {
+            mClass = c;
+            mInstName = instName;
+            mObj = obj;
+        }
+
+        @Override
+        public int hashCode() {
+            return (mClass.getName().hashCode() * 31 + mInstName.hashCode()) * 31;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == null || obj.getClass() != getClass()) {
+                return false;
+            }
+
+            InstanceKey other = (InstanceKey) obj;
+            return (other.mClass == mClass && other.mInstName.equals(mInstName)
+                    && other.mObj == mObj);
+        }
+    }
+
     protected TestContext getTestContext() {
         return mContext;
     }