diff --git a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
index e49c67c..4bb9a53 100644
--- a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
+++ b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
@@ -7,7 +7,6 @@
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
-import org.robolectric.RuntimeEnvironment;
 
 import java.lang.annotation.Annotation;
 import java.lang.annotation.ElementType;
@@ -15,6 +14,10 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 /**
  * Test rule that makes overriding flags in Robolectric tests easier. This rule clears all flags
@@ -52,68 +55,48 @@
         boolean value();
     }
 
-    private boolean ruleInProgress;
-
     @Override
     public Statement apply(Statement base, Description description) {
-        return new Statement() {
-            @Override
-            public void evaluate() throws Throwable {
-                FeatureFlags.initialize(RuntimeEnvironment.application.getApplicationContext());
-                ruleInProgress = true;
-                try {
-                    clearOverrides();
-                    applyAnnotationOverrides(description);
-                    base.evaluate();
-                } finally {
-                    ruleInProgress = false;
-                    clearOverrides();
+        return new MyStatement(base, description);
+    }
+
+    private class MyStatement extends Statement {
+
+        private final Statement mBase;
+        private final Description mDescription;
+
+
+        MyStatement(Statement base, Description description) {
+            mBase = base;
+            mDescription = description;
+        }
+
+        @Override
+        public void evaluate() throws Throwable {
+            Map<String, BaseTogglableFlag> allFlags = FeatureFlags.getTogglableFlags().stream()
+                    .collect(Collectors.toMap(TogglableFlag::getKey, Function.identity()));
+
+            HashMap<BaseTogglableFlag, Boolean> changedValues = new HashMap<>();
+            FlagOverride[] overrides = new FlagOverride[0];
+            try {
+                for (Annotation annotation : mDescription.getAnnotations()) {
+                    if (annotation.annotationType() == FlagOverride.class) {
+                        overrides = new FlagOverride[] { (FlagOverride) annotation };
+                    } else if (annotation.annotationType() == FlagOverrides.class) {
+                        // Note: this branch is hit if the annotation is repeated
+                        overrides = ((FlagOverrides) annotation).value();
+                    }
                 }
-            }
-        };
-    }
-
-    private void override(BaseTogglableFlag flag, boolean newValue) {
-        if (!ruleInProgress) {
-            throw new IllegalStateException(
-                    "Rule isn't in progress. Did you remember to mark it with @Rule?");
-        }
-        flag.setForTests(newValue);
-    }
-
-    private void applyAnnotationOverrides(Description description) {
-        for (Annotation annotation : description.getAnnotations()) {
-            if (annotation.annotationType() == FlagOverride.class) {
-                applyAnnotation((FlagOverride) annotation);
-            } else if (annotation.annotationType() == FlagOverrides.class) {
-                // Note: this branch is hit if the annotation is repeated
-                for (FlagOverride flagOverride : ((FlagOverrides) annotation).value()) {
-                    applyAnnotation(flagOverride);
+                for (FlagOverride override : overrides) {
+                    BaseTogglableFlag flag = allFlags.get(override.key());
+                    changedValues.put(flag, flag.get());
+                    flag.setForTests(override.value());
                 }
+                mBase.evaluate();
+            } finally {
+                // Clear the values
+                changedValues.forEach(BaseTogglableFlag::setForTests);
             }
         }
     }
-
-    private void applyAnnotation(FlagOverride flagOverride) {
-        boolean found = false;
-        for (TogglableFlag flag : FeatureFlags.getTogglableFlags()) {
-            if (flag.getKey().equals(flagOverride.key())) {
-                override(flag, flagOverride.value());
-                found = true;
-                break;
-            }
-        }
-        if (!found) {
-            throw new IllegalStateException("Flag " + flagOverride.key() + " not found");
-        }
-    }
-
-    /**
-     * Resets all flags to their default values.
-     */
-    private void clearOverrides() {
-        for (BaseTogglableFlag flag : FeatureFlags.getTogglableFlags()) {
-            flag.setForTests(flag.getDefaultValue());
-        }
-    }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java b/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
index 096db57..410a077 100644
--- a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
+++ b/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
@@ -1,20 +1,22 @@
 package com.android.launcher3.logging;
 
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
+import org.robolectric.Shadows;
+import org.robolectric.util.Scheduler;
 
 import java.io.File;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.Calendar;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
 /**
  * Tests for {@link FileLog}
  */
@@ -22,9 +24,10 @@
 public class FileLogTest {
 
     private File mTempDir;
+    private boolean mTestActive;
 
     @Before
-    public void setUp() throws Exception {
+    public void setUp() {
         int count = 0;
         do {
             mTempDir = new File(RuntimeEnvironment.application.getCacheDir(),
@@ -32,14 +35,24 @@
         } while (!mTempDir.mkdir());
 
         FileLog.setDir(mTempDir);
+
+        mTestActive = true;
+        Scheduler scheduler = Shadows.shadowOf(FileLog.getHandler().getLooper()).getScheduler();
+        new Thread(() -> {
+            while (mTestActive) {
+                scheduler.advanceToLastPostedRunnable();
+            }
+        }).start();
     }
 
     @After
-    public void tearDown() throws Exception {
+    public void tearDown() {
         // Clear existing logs
         new File(mTempDir, "log-0").delete();
         new File(mTempDir, "log-1").delete();
         mTempDir.delete();
+
+        mTestActive = false;
     }
 
     @Test
@@ -49,12 +62,12 @@
         }
         FileLog.print("Testing", "hoolalala");
         StringWriter writer = new StringWriter();
-        FileLog.flushAll(new PrintWriter(writer));
+        assertTrue(FileLog.flushAll(new PrintWriter(writer)));
         assertTrue(writer.toString().contains("hoolalala"));
 
         FileLog.print("Testing", "abracadabra", new Exception("cat! cat!"));
         writer = new StringWriter();
-        FileLog.flushAll(new PrintWriter(writer));
+        assertTrue(FileLog.flushAll(new PrintWriter(writer)));
         assertTrue(writer.toString().contains("abracadabra"));
         // Exception is also printed
         assertTrue(writer.toString().contains("cat! cat!"));
@@ -70,7 +83,7 @@
         }
         FileLog.print("Testing", "hoolalala");
         StringWriter writer = new StringWriter();
-        FileLog.flushAll(new PrintWriter(writer));
+        assertTrue(FileLog.flushAll(new PrintWriter(writer)));
         assertTrue(writer.toString().contains("hoolalala"));
 
         Calendar threeDaysAgo = Calendar.getInstance();
@@ -80,7 +93,7 @@
 
         FileLog.print("Testing", "abracadabra", new Exception("cat! cat!"));
         writer = new StringWriter();
-        FileLog.flushAll(new PrintWriter(writer));
+        assertTrue(FileLog.flushAll(new PrintWriter(writer)));
         assertTrue(writer.toString().contains("abracadabra"));
         // Exception is also printed
         assertTrue(writer.toString().contains("cat! cat!"));
diff --git a/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index bc936b7..32eb2ec 100644
--- a/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ b/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
@@ -15,17 +15,20 @@
 import android.os.Process;
 import android.os.UserHandle;
 
+import androidx.annotation.NonNull;
+
 import com.android.launcher3.AppFilter;
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
-import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherProvider;
+import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.icons.cache.CachingLogic;
+import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.TestLauncherProvider;
 
@@ -44,8 +47,6 @@
 import java.util.concurrent.Executor;
 import java.util.function.Supplier;
 
-import androidx.annotation.NonNull;
-
 /**
  * Base class for writing tests for Model update tasks.
  */
@@ -79,6 +80,7 @@
         model = mock(LauncherModel.class);
         modelWriter = mock(ModelWriter.class);
 
+        LauncherAppState.INSTANCE.initializeForTesting(appState);
         when(appState.getModel()).thenReturn(model);
         when(model.getWriter(anyBoolean(), anyBoolean())).thenReturn(modelWriter);
         when(model.getCallback()).thenReturn(callbacks);
@@ -216,5 +218,10 @@
         public Bitmap newIcon() {
             return Bitmap.createBitmap(1, 1, Config.ARGB_8888);
         }
+
+        @Override
+        public synchronized BitmapInfo getDefaultIcon(UserHandle user) {
+            return BitmapInfo.fromBitmap(newIcon());
+        }
     }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index 42848f4..81b9043 100644
--- a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -11,7 +11,6 @@
 import com.android.launcher3.WorkspaceItemInfo;
 
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
@@ -41,7 +40,6 @@
     }
 
     @Test
-    @Ignore("This test fails with resource errors") // b/131115553
     public void testCacheUpdate_update_apps() throws Exception {
         // Clear all icons from apps list so that its easy to check what was updated
         for (AppInfo info : allAppsList.data) {
@@ -66,7 +64,6 @@
     }
 
     @Test
-    @Ignore("This test fails with resource errors") // b/131115553
     public void testSessionUpdate_ignores_normal_apps() throws Exception {
         executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app1"));
 
@@ -75,7 +72,6 @@
     }
 
     @Test
-    @Ignore("This test fails with resource errors") // b/131115553
     public void testSessionUpdate_updates_pending_apps() throws Exception {
         executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app3"));
 
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index d70abc2..db94bdb 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -44,7 +44,7 @@
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
 
     // We do not need any synchronization for this variable as its only written on UI thread.
-    private static final MainThreadInitializedObject<LauncherAppState> INSTANCE =
+    public static final MainThreadInitializedObject<LauncherAppState> INSTANCE =
             new MainThreadInitializedObject<>(LauncherAppState::new);
 
     private final Context mContext;
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index e6eced1..c502dd7 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -238,14 +238,7 @@
 
         @Override
         public int hashCode() {
-            int h$ = 1;
-            h$ *= 1000003;
-            h$ ^= key.hashCode();
-            h$ *= 1000003;
-            h$ ^= getDefaultValue() ? 1231 : 1237;
-            h$ *= 1000003;
-            h$ ^= description.hashCode();
-            return h$;
+            return key.hashCode();
         }
     }
 }
diff --git a/src/com/android/launcher3/logging/FileLog.java b/src/com/android/launcher3/logging/FileLog.java
index 923a89b..04cf20a 100644
--- a/src/com/android/launcher3/logging/FileLog.java
+++ b/src/com/android/launcher3/logging/FileLog.java
@@ -8,6 +8,8 @@
 import android.util.Log;
 import android.util.Pair;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.util.IOUtils;
 
 import java.io.BufferedReader;
@@ -88,7 +90,8 @@
         Message.obtain(getHandler(), LogWriterCallback.MSG_WRITE, out).sendToTarget();
     }
 
-    private static Handler getHandler() {
+    @VisibleForTesting
+    static Handler getHandler() {
         synchronized (DATE_FORMAT) {
             if (sHandler == null) {
                 sHandler = new Handler(createAndStartNewLooper("file-logger"),
@@ -102,15 +105,16 @@
      * Blocks until all the pending logs are written to the disk
      * @param out if not null, all the persisted logs are copied to the writer.
      */
-    public static void flushAll(PrintWriter out) throws InterruptedException {
+    public static boolean flushAll(PrintWriter out) throws InterruptedException {
         if (!ENABLED) {
-            return;
+            return false;
         }
         CountDownLatch latch = new CountDownLatch(1);
         Message.obtain(getHandler(), LogWriterCallback.MSG_FLUSH,
                 Pair.create(out, latch)).sendToTarget();
 
         latch.await(2, TimeUnit.SECONDS);
+        return latch.getCount() == 0;
     }
 
     /**