diff --git a/tests/RollbackTest/RollbackTest.xml b/tests/RollbackTest/RollbackTest.xml
index 269cec1..7b85cc8 100644
--- a/tests/RollbackTest/RollbackTest.xml
+++ b/tests/RollbackTest/RollbackTest.xml
@@ -23,6 +23,10 @@
         <option name="run-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\\*' --esa flags &quot;ModuleConfig__versioned_immediate_commit_packages&quot; --esa types &quot;bytes&quot; --esa values &quot;Cm5vdGFwYWNrYWdlOgA=&quot; com.google.android.gms" />
         <option name="teardown-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es action delete --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\*' --esa flag &quot;ModuleConfig__immediate_commit_packages&quot; com.google.android.gms" />
         <option name="teardown-command" value="am broadcast -a 'com.google.android.gms.phenotype.FLAG_OVERRIDE' --es action delete --es package &quot;com.google.android.gms.platformconfigurator&quot; --es user '\*' --esa flag &quot;ModuleConfig__versioned_immediate_commit_packages&quot; com.google.android.gms" />
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="run-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.A" />
+        <option name="teardown-command" value="pm uninstall com.android.cts.install.lib.testapp.B" />
     </target_preparer>
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.tests.rollback" />
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
index 5a92d68..cab8b42 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -75,6 +75,12 @@
     private static final String PROPERTY_ENABLE_ROLLBACK_TIMEOUT_MILLIS =
             "enable_rollback_timeout";
 
+    private static boolean hasRollbackInclude(List<RollbackInfo> rollbacks, String packageName) {
+        return rollbacks.stream().anyMatch(
+                ri -> ri.getPackages().stream().anyMatch(
+                        pri -> packageName.equals(pri.getPackageName())));
+    }
+
     /**
      * Test basic rollbacks.
      */
@@ -113,18 +119,14 @@
             // Uninstall TestApp.A
             Uninstall.packages(TestApp.A);
             assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
-            // TODO: There is currently a race condition between when the app is
-            // uninstalled and when rollback manager deletes the rollback. Fix it
-            // so that's not the case!
             for (int i = 0; i < 5; ++i) {
-                RollbackInfo rollback = getUniqueRollbackInfoForPackage(
-                        rm.getRecentlyCommittedRollbacks(), TestApp.A);
-                if (rollback != null) {
+                if (hasRollbackInclude(rm.getRecentlyCommittedRollbacks(), TestApp.A)) {
                     Log.i(TAG, "Sleeping 1 second to wait for uninstall to take effect.");
                     Thread.sleep(1000);
                 }
             }
 
+            assertThat(hasRollbackInclude(rm.getRecentlyCommittedRollbacks(), TestApp.A)).isFalse();
             // The app should not be available for rollback.
             waitForUnavailableRollback(TestApp.A);