Allow track selection parameters to be set in ExoPlayerAssetLoader.

PiperOrigin-RevId: 701926949

libraries/test_data/src/test/assets/media/dash/multi-track/sample.mpd

@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<MPD xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xmlns="urn:mpeg:dash:schema:mpd:2011"
+	xmlns:xlink="http://www.w3.org/1999/xlink"
+	xsi:schemaLocation="urn:mpeg:DASH:schema:MPD:2011 http://standards.iso.org/ittf/PubliclyAvailableStandards/MPEG-DASH_schema_files/DASH-MPD.xsd"
+	profiles="urn:mpeg:dash:profile:isoff-live:2011"
+	type="static"
+	mediaPresentationDuration="PT1.0S"
+	maxSegmentDuration="PT5.0S"
+	minBufferTime="PT2.0S">
+	<ProgramInformation>
+	</ProgramInformation>
+	<ServiceDescription id="0">
+	</ServiceDescription>
+	<Period id="0" start="PT0.0S">
+		<AdaptationSet id="0" contentType="video" startWithSAP="1" segmentAlignment="true" bitstreamSwitching="true" frameRate="30000/1001" maxWidth="1080" maxHeight="720" par="3:2" lang="und">
+			<Representation id="0" mimeType="video/mp4" codecs="avc1.4d400d" bandwidth="300000" width="320" height="240" sar="9:8">
+				<SegmentTemplate timescale="30000" initialization="init-stream$RepresentationID$.m4s" media="chunk-stream$RepresentationID$-$Number%05d$.m4s" startNumber="1">
+					<SegmentTimeline>
+						<S t="0" d="30030" />
+					</SegmentTimeline>
+				</SegmentTemplate>
+			</Representation>
+			<Representation id="1" mimeType="video/mp4" codecs="avc1.42c01e" bandwidth="3000000" width="640" height="360" sar="27:32">
+				<SegmentTemplate timescale="30000" initialization="init-stream$RepresentationID$.m4s" media="chunk-stream$RepresentationID$-$Number%05d$.m4s" startNumber="1">
+					<SegmentTimeline>
+						<S t="0" d="30030" />
+					</SegmentTimeline>
+				</SegmentTemplate>
+			</Representation>
+		</AdaptationSet>
+	</Period>
+</MPD>

libraries/transformer/build.gradle

@@ -56,6 +56,7 @@ dependencies {
     compileOnly 'com.google.errorprone:error_prone_annotations:' + errorProneVersion
     compileOnly 'org.checkerframework:checker-qual:' + checkerframeworkVersion
     compileOnly 'org.jetbrains.kotlin:kotlin-annotations-jvm:' + kotlinAnnotationsVersion
+    testImplementation project(modulePrefix + 'lib-exoplayer-dash')
     testImplementation project(modulePrefix + 'test-utils-robolectric')
     testImplementation project(modulePrefix + 'test-utils')
     testImplementation project(modulePrefix + 'test-data')

libraries/transformer/src/main/java/androidx/media3/transformer/DefaultAssetLoaderFactory.java

@@ -34,6 +34,7 @@ import androidx.media3.common.util.Util;
 import androidx.media3.datasource.DataSourceBitmapLoader;
 import androidx.media3.datasource.DefaultDataSource;
 import androidx.media3.exoplayer.source.MediaSource;
+import androidx.media3.exoplayer.trackselection.TrackSelector;
 import androidx.media3.transformer.AssetLoader.CompositionSettings;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.Executors;
@@ -56,6 +57,7 @@ public final class DefaultAssetLoaderFactory implements AssetLoader.Factory {
   private final Clock clock;
   @Nullable private final MediaSource.Factory mediaSourceFactory;
   private final BitmapLoader bitmapLoader;
+  @Nullable private final TrackSelector.Factory trackSelectorFactory;
 
   private AssetLoader.@MonotonicNonNull Factory imageAssetLoaderFactory;
   private AssetLoader.@MonotonicNonNull Factory exoPlayerAssetLoaderFactory;
@@ -75,10 +77,12 @@ public final class DefaultAssetLoaderFactory implements AssetLoader.Factory {
    */
   public DefaultAssetLoaderFactory(
       Context context, Codec.DecoderFactory decoderFactory, Clock clock) {
+    // TODO: b/381519379 - deprecate this constructor and replace with a builder.
     this.context = context.getApplicationContext();
     this.decoderFactory = decoderFactory;
     this.clock = clock;
     this.mediaSourceFactory = null;
+    this.trackSelectorFactory = null;
     @Nullable BitmapFactory.Options options = null;
     if (Util.SDK_INT >= 26) {
       options = new BitmapFactory.Options();
@@ -102,11 +106,13 @@ public final class DefaultAssetLoaderFactory implements AssetLoader.Factory {
    * @param bitmapLoader The {@link BitmapLoader} to use to load and decode images.
    */
   public DefaultAssetLoaderFactory(Context context, BitmapLoader bitmapLoader) {
+    // TODO: b/381519379 - deprecate this constructor and replace with a builder.
     this.context = context.getApplicationContext();
     this.bitmapLoader = bitmapLoader;
     decoderFactory = new DefaultDecoderFactory.Builder(context).build();
     clock = Clock.DEFAULT;
     mediaSourceFactory = null;
+    trackSelectorFactory = null;
   }
 
   /**
@@ -127,11 +133,43 @@ public final class DefaultAssetLoaderFactory implements AssetLoader.Factory {
       Clock clock,
       @Nullable MediaSource.Factory mediaSourceFactory,
       BitmapLoader bitmapLoader) {
+    // TODO: b/381519379 - deprecate this constructor and replace with a builder.
     this.context = context.getApplicationContext();
     this.decoderFactory = decoderFactory;
     this.clock = clock;
     this.mediaSourceFactory = mediaSourceFactory;
     this.bitmapLoader = bitmapLoader;
+    this.trackSelectorFactory = null;
+  }
+
+  /**
+   * Creates an instance.
+   *
+   * @param context The {@link Context}.
+   * @param decoderFactory The {@link Codec.DecoderFactory} to use to decode the samples (if
+   *     necessary).
+   * @param clock The {@link Clock} to use. It should always be {@link Clock#DEFAULT}, except for
+   *     testing.
+   * @param mediaSourceFactory The {@link MediaSource.Factory} to use to retrieve the samples to
+   *     transform when an {@link ExoPlayerAssetLoader} is used.
+   * @param bitmapLoader The {@link BitmapLoader} to use to load and decode images.
+   * @param trackSelectorFactory The {@link TrackSelector.Factory} to use when selecting the track
+   *     to transform.
+   */
+  public DefaultAssetLoaderFactory(
+      Context context,
+      Codec.DecoderFactory decoderFactory,
+      Clock clock,
+      @Nullable MediaSource.Factory mediaSourceFactory,
+      BitmapLoader bitmapLoader,
+      TrackSelector.Factory trackSelectorFactory) {
+    // TODO: b/381519379 - deprecate this constructor and replace with a builder.
+    this.context = context.getApplicationContext();
+    this.decoderFactory = decoderFactory;
+    this.clock = clock;
+    this.mediaSourceFactory = mediaSourceFactory;
+    this.bitmapLoader = bitmapLoader;
+    this.trackSelectorFactory = trackSelectorFactory;
   }
 
   @Override
@@ -157,9 +195,8 @@ public final class DefaultAssetLoaderFactory implements AssetLoader.Factory {
     }
     if (exoPlayerAssetLoaderFactory == null) {
       exoPlayerAssetLoaderFactory =
-          mediaSourceFactory != null
+          new ExoPlayerAssetLoader.Factory(
-              ? new ExoPlayerAssetLoader.Factory(context, decoderFactory, clock, mediaSourceFactory)
+              context, decoderFactory, clock, mediaSourceFactory, trackSelectorFactory);
-              : new ExoPlayerAssetLoader.Factory(context, decoderFactory, clock);
     }
     return exoPlayerAssetLoaderFactory.createAssetLoader(
         editedMediaItem, looper, listener, compositionSettings);

libraries/transformer/src/main/java/androidx/media3/transformer/ExoPlayerAssetLoader.java

@@ -53,6 +53,7 @@ import androidx.media3.exoplayer.source.DefaultMediaSourceFactory;
 import androidx.media3.exoplayer.source.MediaSource;
 import androidx.media3.exoplayer.text.TextOutput;
 import androidx.media3.exoplayer.trackselection.DefaultTrackSelector;
+import androidx.media3.exoplayer.trackselection.TrackSelector;
 import androidx.media3.exoplayer.video.VideoRendererEventListener;
 import androidx.media3.extractor.DefaultExtractorsFactory;
 import androidx.media3.extractor.mp4.Mp4Extractor;
@@ -70,6 +71,7 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
     private final Codec.DecoderFactory decoderFactory;
     private final Clock clock;
     @Nullable private final MediaSource.Factory mediaSourceFactory;
+    @Nullable private final TrackSelector.Factory trackSelectorFactory;
 
     /**
      * Creates an instance using a {@link DefaultMediaSourceFactory}.
@@ -81,10 +83,13 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
      *     testing.
      */
     public Factory(Context context, Codec.DecoderFactory decoderFactory, Clock clock) {
-      this.context = context;
+      // TODO: b/381519379 - deprecate this constructor and replace with a builder.
-      this.decoderFactory = decoderFactory;
+      this(
-      this.clock = clock;
+          context,
-      this.mediaSourceFactory = null;
+          decoderFactory,
+          clock,
+          /* mediaSourceFactory= */ null,
+          /* trackSelectorFactory= */ null);
     }
 
     /**
@@ -103,10 +108,35 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
         Codec.DecoderFactory decoderFactory,
         Clock clock,
         MediaSource.Factory mediaSourceFactory) {
+      // TODO: b/381519379 - deprecate this constructor and replace with a builder.
+      this(context, decoderFactory, clock, mediaSourceFactory, /* trackSelectorFactory= */ null);
+    }
+
+    /**
+     * Creates an instance.
+     *
+     * @param context The {@link Context}.
+     * @param decoderFactory The {@link Codec.DecoderFactory} to use to decode the samples (if
+     *     necessary).
+     * @param clock The {@link Clock} to use. It should always be {@link Clock#DEFAULT}, except for
+     *     testing.
+     * @param mediaSourceFactory The {@link MediaSource.Factory} to use to retrieve the samples to
+     *     transform.
+     * @param trackSelectorFactory The {@link TrackSelector.Factory} to use when selecting the track
+     *     to transform.
+     */
+    public Factory(
+        Context context,
+        Codec.DecoderFactory decoderFactory,
+        Clock clock,
+        @Nullable MediaSource.Factory mediaSourceFactory,
+        @Nullable TrackSelector.Factory trackSelectorFactory) {
+      // TODO: b/381519379 - deprecate this constructor and replace with a builder.
       this.context = context;
       this.decoderFactory = decoderFactory;
       this.clock = clock;
       this.mediaSourceFactory = mediaSourceFactory;
+      this.trackSelectorFactory = trackSelectorFactory;
     }
 
     @Override
@@ -123,6 +153,20 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
         }
         mediaSourceFactory = new DefaultMediaSourceFactory(context, defaultExtractorsFactory);
       }
+      TrackSelector.Factory trackSelectorFactory = this.trackSelectorFactory;
+      if (trackSelectorFactory == null) {
+        DefaultTrackSelector.Parameters defaultTrackSelectorParameters =
+            new DefaultTrackSelector.Parameters.Builder(context)
+                .setForceHighestSupportedBitrate(true)
+                .setConstrainAudioChannelCountToDeviceCapabilities(false)
+                .build();
+        trackSelectorFactory =
+            context -> {
+              DefaultTrackSelector trackSelector = new DefaultTrackSelector(context);
+              trackSelector.setParameters(defaultTrackSelectorParameters);
+              return trackSelector;
+            };
+      }
       return new ExoPlayerAssetLoader(
           context,
           editedMediaItem,
@@ -131,7 +175,8 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
           compositionSettings.hdrMode,
           looper,
           listener,
-          clock);
+          clock,
+          trackSelectorFactory);
     }
   }
 
@@ -158,17 +203,13 @@ public final class ExoPlayerAssetLoader implements AssetLoader {
       @Composition.HdrMode int hdrMode,
       Looper looper,
       Listener listener,
-      Clock clock) {
+      Clock clock,
+      TrackSelector.Factory trackSelectorFactory) {
     this.context = context;
     this.editedMediaItem = editedMediaItem;
     this.decoderFactory = new CapturingDecoderFactory(decoderFactory);
 
-    DefaultTrackSelector trackSelector = new DefaultTrackSelector(context);
+    TrackSelector trackSelector = trackSelectorFactory.createTrackSelector(context);
-    trackSelector.setParameters(
-        new DefaultTrackSelector.Parameters.Builder(context)
-            .setForceHighestSupportedBitrate(true)
-            .setConstrainAudioChannelCountToDeviceCapabilities(false)
-            .build());
     // Arbitrarily decrease buffers for playback so that samples start being sent earlier to the
     // exporters (rebuffers are less problematic for the export use case).
     DefaultLoadControl loadControl =

libraries/transformer/src/test/java/androidx/media3/transformer/ExoPlayerAssetLoaderTest.java

@@ -27,6 +27,8 @@ import androidx.media3.common.Format;
 import androidx.media3.common.MediaItem;
 import androidx.media3.common.util.Clock;
 import androidx.media3.decoder.DecoderInputBuffer;
+import androidx.media3.exoplayer.trackselection.DefaultTrackSelector;
+import androidx.media3.exoplayer.trackselection.TrackSelector;
 import androidx.media3.transformer.AssetLoader.CompositionSettings;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -41,91 +43,27 @@ import org.robolectric.shadows.ShadowSystemClock;
 @RunWith(AndroidJUnit4.class)
 public class ExoPlayerAssetLoaderTest {
 
+  private static final String SINGLE_TRACK_URI = "asset:///media/mp4/sample.mp4";
+  // Contains two representations of asset:///assets/media/dash/ttml-in-mp4/sample.video.mp4
+  // one at 360p and the other at 240p.
+  private static final String MULTI_TRACK_URI = "asset:///media/dash/multi-track/sample.mpd";
+
   @Test
   public void exoPlayerAssetLoader_callsListenerCallbacksInRightOrder() throws Exception {
-    AtomicReference<Exception> exceptionRef = new AtomicReference<>();
+    AtomicReference<Exception> exception = new AtomicReference<>();
     AtomicBoolean isAudioOutputFormatSet = new AtomicBoolean();
     AtomicBoolean isVideoOutputFormatSet = new AtomicBoolean();
-
     AssetLoader.Listener listener =
-        new AssetLoader.Listener() {
+        getAssetLoaderListener(
-
+            exception,
-          private volatile boolean isDurationSet;
+            isAudioOutputFormatSet,
-          private volatile boolean isTrackCountSet;
+            isVideoOutputFormatSet,
-          private volatile boolean isAudioTrackAdded;
+            /* expectedOutputResolutionHeight= */ null);
-          private volatile boolean isVideoTrackAdded;
-
-          @Override
-          public void onDurationUs(long durationUs) {
-            // Sleep to increase the chances of the test failing.
-            sleep();
-            isDurationSet = true;
-          }
-
-          @Override
-          public void onTrackCount(int trackCount) {
-            // Sleep to increase the chances of the test failing.
-            sleep();
-            isTrackCountSet = true;
-          }
-
-          @Override
-          public boolean onTrackAdded(
-              Format inputFormat, @AssetLoader.SupportedOutputTypes int supportedOutputTypes) {
-            if (!isDurationSet) {
-              exceptionRef.set(
-                  new IllegalStateException("onTrackAdded() called before onDurationUs()"));
-            } else if (!isTrackCountSet) {
-              exceptionRef.set(
-                  new IllegalStateException("onTrackAdded() called before onTrackCount()"));
-            }
-            sleep();
-            @C.TrackType int trackType = getProcessedTrackType(inputFormat.sampleMimeType);
-            if (trackType == C.TRACK_TYPE_AUDIO) {
-              isAudioTrackAdded = true;
-            } else if (trackType == C.TRACK_TYPE_VIDEO) {
-              isVideoTrackAdded = true;
-            }
-            return false;
-          }
-
-          @Override
-          public SampleConsumer onOutputFormat(Format format) {
-            @C.TrackType int trackType = getProcessedTrackType(format.sampleMimeType);
-            boolean isAudio = trackType == C.TRACK_TYPE_AUDIO;
-            boolean isVideo = trackType == C.TRACK_TYPE_VIDEO;
-
-            boolean isTrackAdded = (isAudio && isAudioTrackAdded) || (isVideo && isVideoTrackAdded);
-            if (!isTrackAdded) {
-              exceptionRef.set(
-                  new IllegalStateException("onOutputFormat() called before onTrackAdded()"));
-            }
-            if (isAudio) {
-              isAudioOutputFormatSet.set(true);
-            } else if (isVideo) {
-              isVideoOutputFormatSet.set(true);
-            }
-            return new FakeSampleConsumer();
-          }
-
-          @Override
-          public void onError(ExportException e) {
-            exceptionRef.set(e);
-          }
-
-          private void sleep() {
-            try {
-              Thread.sleep(10);
-            } catch (InterruptedException e) {
-              Thread.currentThread().interrupt();
-              exceptionRef.set(e);
-            }
-          }
-        };
     // Use default clock so that messages sent on different threads are not always executed in the
     // order in which they are received.
     Clock clock = Clock.DEFAULT;
-    AssetLoader assetLoader = getAssetLoader(listener, clock);
+    AssetLoader assetLoader =
+        getAssetLoader(listener, clock, SINGLE_TRACK_URI, /* trackSelectorFactory= */ null);
 
     assetLoader.start();
     runLooperUntil(
@@ -133,18 +71,109 @@ public class ExoPlayerAssetLoaderTest {
         () -> {
           ShadowSystemClock.advanceBy(Duration.ofMillis(10));
           return (isAudioOutputFormatSet.get() && isVideoOutputFormatSet.get())
-              || exceptionRef.get() != null;
+              || exception.get() != null;
         });
 
-    assertThat(exceptionRef.get()).isNull();
+    assertThat(exception.get()).isNull();
   }
 
-  private static AssetLoader getAssetLoader(AssetLoader.Listener listener, Clock clock) {
+  @Test
+  public void exoPlayerAssetLoader_withMaxVideoSize_loadsLowResolutionTrack() throws Exception {
+    AtomicReference<Exception> exception = new AtomicReference<>();
+    AtomicBoolean isAudioOutputFormatSet = new AtomicBoolean();
+    AtomicBoolean isVideoOutputFormatSet = new AtomicBoolean();
+    int expectedOutputResolutionHeight = 240;
+    AssetLoader.Listener listener =
+        getAssetLoaderListener(
+            exception,
+            isAudioOutputFormatSet,
+            isVideoOutputFormatSet,
+            expectedOutputResolutionHeight);
+    DefaultTrackSelector.Parameters trackSelectorParameters =
+        new DefaultTrackSelector.Parameters.Builder(ApplicationProvider.getApplicationContext())
+            .setMaxVideoSize(
+                /* maxVideoWidth= */ Integer.MAX_VALUE,
+                /* maxVideoHeight= */ expectedOutputResolutionHeight)
+            .setForceHighestSupportedBitrate(true)
+            .setConstrainAudioChannelCountToDeviceCapabilities(false)
+            .build();
+    TrackSelector.Factory trackSelectorFactory =
+        context -> {
+          DefaultTrackSelector trackSelector = new DefaultTrackSelector(context);
+          trackSelector.setParameters(trackSelectorParameters);
+          return trackSelector;
+        };
+    // Use default clock so that messages sent on different threads are not always executed in the
+    // order in which they are received.
+    Clock clock = Clock.DEFAULT;
+    AssetLoader assetLoader =
+        getAssetLoader(listener, clock, MULTI_TRACK_URI, trackSelectorFactory);
+
+    assetLoader.start();
+    runLooperUntil(
+        Looper.myLooper(),
+        () -> {
+          ShadowSystemClock.advanceBy(Duration.ofMillis(10));
+          return isVideoOutputFormatSet.get() || exception.get() != null;
+        });
+
+    // The resolution of the selected track is checked against expectedOutputResolutionHeight in
+    // listener.onOutputFormat.
+    assertThat(exception.get()).isNull();
+  }
+
+  @Test
+  public void exoPlayerAssetLoader_withNoMaxVideoSize_loadsHighResolutionTrack() throws Exception {
+    AtomicReference<Exception> exception = new AtomicReference<>();
+    AtomicBoolean isAudioOutputFormatSet = new AtomicBoolean();
+    AtomicBoolean isVideoOutputFormatSet = new AtomicBoolean();
+    int expectedOutputResolutionHeight = 360;
+    AssetLoader.Listener listener =
+        getAssetLoaderListener(
+            exception,
+            isAudioOutputFormatSet,
+            isVideoOutputFormatSet,
+            expectedOutputResolutionHeight);
+    DefaultTrackSelector.Parameters trackSelectorParameters =
+        new DefaultTrackSelector.Parameters.Builder(ApplicationProvider.getApplicationContext())
+            .setForceHighestSupportedBitrate(true)
+            .setConstrainAudioChannelCountToDeviceCapabilities(false)
+            .build();
+    TrackSelector.Factory trackSelectorFactory =
+        context -> {
+          DefaultTrackSelector trackSelector = new DefaultTrackSelector(context);
+          trackSelector.setParameters(trackSelectorParameters);
+          return trackSelector;
+        };
+    // Use default clock so that messages sent on different threads are not always executed in the
+    // order in which they are received.
+    Clock clock = Clock.DEFAULT;
+    AssetLoader assetLoader =
+        getAssetLoader(listener, clock, MULTI_TRACK_URI, trackSelectorFactory);
+
+    assetLoader.start();
+    runLooperUntil(
+        Looper.myLooper(),
+        () -> {
+          ShadowSystemClock.advanceBy(Duration.ofMillis(10));
+          return isVideoOutputFormatSet.get() || exception.get() != null;
+        });
+
+    // The resolution of the selected track is checked against expectedOutputResolutionHeight in
+    // listener.onOutputFormat.
+    assertThat(exception.get()).isNull();
+  }
+
+  private static AssetLoader getAssetLoader(
+      AssetLoader.Listener listener,
+      Clock clock,
+      String uri,
+      @Nullable TrackSelector.Factory trackSelectorFactory) {
     Context context = ApplicationProvider.getApplicationContext();
     Codec.DecoderFactory decoderFactory = new DefaultDecoderFactory.Builder(context).build();
-    EditedMediaItem editedMediaItem =
+    EditedMediaItem editedMediaItem = new EditedMediaItem.Builder(MediaItem.fromUri(uri)).build();
-        new EditedMediaItem.Builder(MediaItem.fromUri("asset:///media/mp4/sample.mp4")).build();
+    return new ExoPlayerAssetLoader.Factory(
-    return new ExoPlayerAssetLoader.Factory(context, decoderFactory, clock)
+            context, decoderFactory, clock, /* mediaSourceFactory= */ null, trackSelectorFactory)
         .createAssetLoader(
             editedMediaItem,
             Looper.myLooper(),
@@ -153,6 +182,95 @@ public class ExoPlayerAssetLoaderTest {
                 Composition.HDR_MODE_KEEP_HDR, /* retainHdrFromUltraHdrImage= */ false));
   }
 
+  private static AssetLoader.Listener getAssetLoaderListener(
+      AtomicReference<Exception> exceptionRef,
+      AtomicBoolean isAudioOutputFormatSet,
+      AtomicBoolean isVideoOutputFormatSet,
+      @Nullable Integer expectedOutputResolutionHeight) {
+    return new AssetLoader.Listener() {
+
+      private volatile boolean isDurationSet;
+      private volatile boolean isTrackCountSet;
+      private volatile boolean isAudioTrackAdded;
+      private volatile boolean isVideoTrackAdded;
+
+      @Override
+      public void onDurationUs(long durationUs) {
+        // Sleep to increase the chances of the test failing.
+        sleep();
+        isDurationSet = true;
+      }
+
+      @Override
+      public void onTrackCount(int trackCount) {
+        // Sleep to increase the chances of the test failing.
+        sleep();
+        isTrackCountSet = true;
+      }
+
+      @Override
+      public boolean onTrackAdded(
+          Format inputFormat, @AssetLoader.SupportedOutputTypes int supportedOutputTypes) {
+        if (!isDurationSet) {
+          exceptionRef.set(
+              new IllegalStateException("onTrackAdded() called before onDurationUs()"));
+        } else if (!isTrackCountSet) {
+          exceptionRef.set(
+              new IllegalStateException("onTrackAdded() called before onTrackCount()"));
+        }
+        sleep();
+        @C.TrackType int trackType = getProcessedTrackType(inputFormat.sampleMimeType);
+        if (trackType == C.TRACK_TYPE_AUDIO) {
+          isAudioTrackAdded = true;
+        } else if (trackType == C.TRACK_TYPE_VIDEO) {
+          isVideoTrackAdded = true;
+        }
+        return false;
+      }
+
+      @Override
+      public SampleConsumer onOutputFormat(Format format) {
+        @C.TrackType int trackType = getProcessedTrackType(format.sampleMimeType);
+        boolean isAudio = trackType == C.TRACK_TYPE_AUDIO;
+        boolean isVideo = trackType == C.TRACK_TYPE_VIDEO;
+
+        boolean isTrackAdded = (isAudio && isAudioTrackAdded) || (isVideo && isVideoTrackAdded);
+        if (!isTrackAdded) {
+          exceptionRef.set(
+              new IllegalStateException("onOutputFormat() called before onTrackAdded()"));
+        }
+        if (isAudio) {
+          isAudioOutputFormatSet.set(true);
+        } else if (isVideo) {
+          if (expectedOutputResolutionHeight != null
+              && expectedOutputResolutionHeight != format.height) {
+            exceptionRef.set(
+                new IllegalStateException(
+                    String.format(
+                        "Expected output height %s but received output height %s.",
+                        expectedOutputResolutionHeight, format.height)));
+          }
+          isVideoOutputFormatSet.set(true);
+        }
+        return new FakeSampleConsumer();
+      }
+
+      @Override
+      public void onError(ExportException e) {
+        exceptionRef.set(e);
+      }
+
+      private void sleep() {
+        try {
+          Thread.sleep(10);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          exceptionRef.set(e);
+        }
+      }
+    };
+  }
+
   private static final class FakeSampleConsumer implements SampleConsumer {
 
     @Nullable