Merge "TelephonyBackupAgentTest: Clean up after ourselves."
diff --git a/src/com/android/providers/telephony/TelephonyBackupAgent.java b/src/com/android/providers/telephony/TelephonyBackupAgent.java
index ed19f27..ab11957 100644
--- a/src/com/android/providers/telephony/TelephonyBackupAgent.java
+++ b/src/com/android/providers/telephony/TelephonyBackupAgent.java
@@ -702,9 +702,17 @@
         }
     };
 
+    /**
+     * Sets a temporary {@code SmsProviderQuery} for testing; note that this method
+     * is not thread safe.
+     *
+     * @return the previous {@code SmsProviderQuery}
+     */
     @VisibleForTesting
-    public void setSmsProviderQuery(SmsProviderQuery smsProviderQuery) {
+    public SmsProviderQuery getAndSetSmsProviderQuery(SmsProviderQuery smsProviderQuery) {
+        SmsProviderQuery result = mSmsProviderQuery;
         mSmsProviderQuery = smsProviderQuery;
+        return result;
     }
 
     private boolean doesMmsExist(Mms mms) {
diff --git a/tests/src/com/android/providers/telephony/TelephonyBackupAgentTest.java b/tests/src/com/android/providers/telephony/TelephonyBackupAgentTest.java
index 00bb15e..f8cc8a6 100644
--- a/tests/src/com/android/providers/telephony/TelephonyBackupAgentTest.java
+++ b/tests/src/com/android/providers/telephony/TelephonyBackupAgentTest.java
@@ -635,12 +635,16 @@
                         return false;
                     }
         };
-        mTelephonyBackupAgent.setSmsProviderQuery(smsProviderQuery);
-
-        mTelephonyBackupAgent.putSmsMessagesToProvider(jsonReader);
-        // the "- 1" is due to exception thrown for one of the messages
-        assertEquals(mSmsRows.length - 1, smsProvider.getRowsAdded());
-        assertEquals(mThreadProvider.mIsThreadArchived, mThreadProvider.mUpdateThreadsArchived);
+        TelephonyBackupAgent.SmsProviderQuery previousQuery =
+                mTelephonyBackupAgent.getAndSetSmsProviderQuery(smsProviderQuery);
+        try {
+            mTelephonyBackupAgent.putSmsMessagesToProvider(jsonReader);
+            // the "- 1" is due to exception thrown for one of the messages
+            assertEquals(mSmsRows.length - 1, smsProvider.getRowsAdded());
+            assertEquals(mThreadProvider.mIsThreadArchived, mThreadProvider.mUpdateThreadsArchived);
+        } finally {
+            mTelephonyBackupAgent.getAndSetSmsProviderQuery(previousQuery);
+        }
     }
 
     /**