1 files changed, 10 insertions, 1 deletions
diff --git a/src/com/android/dialer/compat/FilteredNumberCompat.java b/src/com/android/dialer/compat/FilteredNumberCompat.java
index 91563dc0d..008782d86 100644
--- a/src/com/android/dialer/compat/FilteredNumberCompat.java
+++ b/src/com/android/dialer/compat/FilteredNumberCompat.java
@@ -16,6 +16,7 @@
 
 package com.android.dialer.compat;
 
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 
 import android.app.FragmentManager;
@@ -67,6 +68,8 @@ public class FilteredNumberCompat {
 
     private static Boolean isEnabledForTest;
 
+    private static Context contextForTest;
+
     /**
      * @return The column name for ID in the filtered number database.
      */
@@ -154,7 +157,8 @@ public class FilteredNumberCompat {
      */
     @NeededForTesting
     public static void setHasMigratedToNewBlocking(boolean hasMigrated) {
-        PreferenceManager.getDefaultSharedPreferences(DialerApplication.getContext()).edit()
+        PreferenceManager.getDefaultSharedPreferences(
+                MoreObjects.firstNonNull(contextForTest, DialerApplication.getContext())).edit()
                 .putBoolean(HAS_MIGRATED_TO_NEW_BLOCKING_KEY, hasMigrated).apply();
     }
 
@@ -163,6 +167,11 @@ public class FilteredNumberCompat {
         isEnabledForTest = isEnabled;
     }
 
+    @NeededForTesting
+    public static void setContextForTest(Context context) {
+        contextForTest = context;
+    }
+
     /**
      * Gets the content {@link Uri} for number filtering.
      *