Fix flaky DownloadManagerDashTest

PiperOrigin-RevId: 244170179

library/core/src/test/java/com/google/android/exoplayer2/offline/DownloadManagerTest.java

@@ -552,7 +552,7 @@ public class DownloadManagerTest {
     }
   }
 
-  private void runOnMainThread(final TestRunnable r) {
+  private void runOnMainThread(TestRunnable r) {
     dummyMainThread.runTestOnMainThread(r);
   }
 

library/dash/src/test/java/com/google/android/exoplayer2/source/dash/offline/DownloadManagerDashTest.java

@@ -34,6 +34,7 @@ import com.google.android.exoplayer2.offline.DownloaderConstructorHelper;
 import com.google.android.exoplayer2.offline.StreamKey;
 import com.google.android.exoplayer2.scheduler.Requirements;
 import com.google.android.exoplayer2.testutil.DummyMainThread;
+import com.google.android.exoplayer2.testutil.DummyMainThread.TestRunnable;
 import com.google.android.exoplayer2.testutil.FakeDataSet;
 import com.google.android.exoplayer2.testutil.FakeDataSource;
 import com.google.android.exoplayer2.testutil.RobolectricUtil;
@@ -100,8 +101,8 @@ public class DownloadManagerDashTest {
   }
 
   @After
-  public void tearDown() throws Exception {
-    downloadManager.release();
+  public void tearDown() {
+    runOnMainThread(() -> downloadManager.release());
     Util.recursiveDelete(tempFolder);
     dummyMainThread.release();
   }
@@ -129,10 +130,11 @@ public class DownloadManagerDashTest {
 
     // Run DM accessing code on UI/main thread as it should be. Also not to block handling of loaded
     // actions.
-    dummyMainThread.runOnMainThread(
+    runOnMainThread(
         () -> {
           // Setup an Action and immediately release the DM.
-          handleDownloadRequest(fakeStreamKey1, fakeStreamKey2);
+          DownloadRequest request = getDownloadRequest(fakeStreamKey1, fakeStreamKey2);
+          downloadManager.addDownload(request);
           downloadManager.release();
         });
 
@@ -229,25 +231,28 @@ public class DownloadManagerDashTest {
   }
 
   private void handleDownloadRequest(StreamKey... keys) {
+    DownloadRequest request = getDownloadRequest(keys);
+    runOnMainThread(() -> downloadManager.addDownload(request));
+  }
+
+  private DownloadRequest getDownloadRequest(StreamKey... keys) {
     ArrayList<StreamKey> keysList = new ArrayList<>();
     Collections.addAll(keysList, keys);
-    DownloadRequest action =
-        new DownloadRequest(
-            TEST_ID,
-            DownloadRequest.TYPE_DASH,
-            TEST_MPD_URI,
-            keysList,
-            /* customCacheKey= */ null,
-            null);
-    downloadManager.addDownload(action);
+    return new DownloadRequest(
+        TEST_ID,
+        DownloadRequest.TYPE_DASH,
+        TEST_MPD_URI,
+        keysList,
+        /* customCacheKey= */ null,
+        null);
   }
 
   private void handleRemoveAction() {
-    downloadManager.removeDownload(TEST_ID);
+    runOnMainThread(() -> downloadManager.removeDownload(TEST_ID));
   }
 
   private void createDownloadManager() {
-    dummyMainThread.runTestOnMainThread(
+    runOnMainThread(
         () -> {
           Factory fakeDataSourceFactory = new FakeDataSource.Factory().setFakeDataSet(fakeDataSet);
           downloadManager =
@@ -261,9 +266,13 @@ public class DownloadManagerDashTest {
                   new Requirements(0));
 
           downloadManagerListener =
-              new TestDownloadManagerListener(downloadManager, dummyMainThread);
+              new TestDownloadManagerListener(
+                  downloadManager, dummyMainThread, /* timeout= */ 3000);
           downloadManager.startDownloads();
         });
   }
 
+  private void runOnMainThread(TestRunnable r) {
+    dummyMainThread.runTestOnMainThread(r);
+  }
 }

testutils_robolectric/src/main/java/com/google/android/exoplayer2/testutil/TestDownloadManagerListener.java

@@ -40,14 +40,21 @@ public final class TestDownloadManagerListener implements DownloadManager.Listen
   private final DummyMainThread dummyMainThread;
   private final HashMap<String, ArrayBlockingQueue<Integer>> downloadStates;
   private final ConditionVariable initializedCondition;
+  private final int timeout;
 
   private CountDownLatch downloadFinishedCondition;
   @Download.FailureReason private int failureReason;
 
   public TestDownloadManagerListener(
       DownloadManager downloadManager, DummyMainThread dummyMainThread) {
+    this(downloadManager, dummyMainThread, TIMEOUT);
+  }
+
+  public TestDownloadManagerListener(
+      DownloadManager downloadManager, DummyMainThread dummyMainThread, int timeout) {
     this.downloadManager = downloadManager;
     this.dummyMainThread = dummyMainThread;
+    this.timeout = timeout;
     downloadStates = new HashMap<>();
     initializedCondition = new ConditionVariable();
     downloadManager.addListener(this);
@@ -110,7 +117,7 @@ public final class TestDownloadManagerListener implements DownloadManager.Listen
             downloadFinishedCondition.countDown();
           }
         });
-    assertThat(downloadFinishedCondition.await(TIMEOUT, TimeUnit.MILLISECONDS)).isTrue();
+    assertThat(downloadFinishedCondition.await(timeout, TimeUnit.MILLISECONDS)).isTrue();
   }
 
   private ArrayBlockingQueue<Integer> getStateQueue(String taskId) {