diff --git a/testutils/src/main/java/com/google/android/exoplayer2/testutil/DummyMainThread.java b/testutils/src/main/java/com/google/android/exoplayer2/testutil/DummyMainThread.java
index 982fcdf558..c10cadbb04 100644
--- a/testutils/src/main/java/com/google/android/exoplayer2/testutil/DummyMainThread.java
+++ b/testutils/src/main/java/com/google/android/exoplayer2/testutil/DummyMainThread.java
@@ -53,7 +53,7 @@ public final class DummyMainThread {
    * Runs the provided {@link Runnable} on the main thread, blocking until execution completes or
    * until timeout milliseconds have passed.
    *
-   * @param timeoutMs the maximum time to wait in milliseconds.
+   * @param timeoutMs The maximum time to wait in milliseconds.
    * @param runnable The {@link Runnable} to run.
    */
   public void runOnMainThread(int timeoutMs, final Runnable runnable) {
diff --git a/testutils_robolectric/src/main/java/com/google/android/exoplayer2/testutil/MediaPeriodAsserts.java b/testutils_robolectric/src/main/java/com/google/android/exoplayer2/testutil/MediaPeriodAsserts.java
index b4137a41de..e99e2f45bb 100644
--- a/testutils_robolectric/src/main/java/com/google/android/exoplayer2/testutil/MediaPeriodAsserts.java
+++ b/testutils_robolectric/src/main/java/com/google/android/exoplayer2/testutil/MediaPeriodAsserts.java
@@ -187,13 +187,14 @@ public final class MediaPeriodAsserts {
   private static TrackGroupArray getTrackGroups(MediaPeriod mediaPeriod) {
     AtomicReference<TrackGroupArray> trackGroupArray = new AtomicReference<>(null);
     DummyMainThread dummyMainThread = new DummyMainThread();
+    ConditionVariable preparedCondition = new ConditionVariable();
     dummyMainThread.runOnMainThread(
         () -> {
-          ConditionVariable preparedCondition = new ConditionVariable();
           mediaPeriod.prepare(
               new Callback() {
                 @Override
                 public void onPrepared(MediaPeriod mediaPeriod) {
+                  trackGroupArray.set(mediaPeriod.getTrackGroups());
                   preparedCondition.open();
                 }
 
@@ -203,13 +204,12 @@ public final class MediaPeriodAsserts {
                 }
               },
               /* positionUs= */ 0);
-          try {
-            preparedCondition.block();
-          } catch (InterruptedException e) {
-            // Ignore.
-          }
-          trackGroupArray.set(mediaPeriod.getTrackGroups());
         });
+    try {
+      preparedCondition.block();
+    } catch (InterruptedException e) {
+      // Ignore.
+    }
     dummyMainThread.release();
     return trackGroupArray.get();
   }