Add frame count tests for preview

This is to ensure prewarming doesn't introduce any regression

PiperOrigin-RevId: 657559693

libraries/transformer/src/androidTest/java/androidx/media3/transformer/AndroidTestUtil.java

@@ -58,6 +58,7 @@ import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.json.JSONException;
 import org.json.JSONObject;
 import org.junit.AssumptionViolatedException;
@@ -77,13 +78,13 @@ public final class AndroidTestUtil {
   public static final class AssetInfo {
     private static final class Builder {
       private final String uri;
-      @Nullable private Format videoFormat;
+      private @MonotonicNonNull Format videoFormat;
       private int videoFrameCount;
       private long videoDurationUs;
+      private @MonotonicNonNull ImmutableList<Long> videoTimestampsUs;
 
       public Builder(String uri) {
         this.uri = uri;
-        videoFormat = null;
         videoFrameCount = C.LENGTH_UNSET;
         videoDurationUs = C.TIME_UNSET;
       }
@@ -111,9 +112,21 @@ public final class AndroidTestUtil {
         return this;
       }
 
+      /** See {@link AssetInfo#videoTimestampsUs}. */
+      @CanIgnoreReturnValue
+      public Builder setVideoTimestampsUs(ImmutableList<Long> videoTimestampsUs) {
+        this.videoTimestampsUs = videoTimestampsUs;
+        return this;
+      }
+
       /** Creates an {@link AssetInfo}. */
       public AssetInfo build() {
-        return new AssetInfo(uri, videoFormat, videoDurationUs, videoFrameCount);
+        if (videoTimestampsUs != null) {
+          checkState(
+              videoFrameCount == C.LENGTH_UNSET || videoFrameCount == videoTimestampsUs.size());
+          videoFrameCount = videoTimestampsUs.size();
+        }
+        return new AssetInfo(uri, videoFormat, videoDurationUs, videoFrameCount, videoTimestampsUs);
       }
     }
 
@@ -129,12 +142,20 @@ public final class AndroidTestUtil {
     /** Video frame count, or {@link C#LENGTH_UNSET}. */
     public final int videoFrameCount;
 
+    /** Video frame timestamps in microseconds, or {@code null}. */
+    @Nullable public final ImmutableList<Long> videoTimestampsUs;
+
     private AssetInfo(
-        String uri, @Nullable Format videoFormat, long videoDurationUs, int videoFrameCount) {
+        String uri,
+        @Nullable Format videoFormat,
+        long videoDurationUs,
+        int videoFrameCount,
+        @Nullable ImmutableList<Long> videoTimestampsUs) {
       this.uri = uri;
       this.videoFormat = videoFormat;
       this.videoDurationUs = videoDurationUs;
       this.videoFrameCount = videoFrameCount;
+      this.videoTimestampsUs = videoTimestampsUs;
     }
 
     @Override
@@ -247,7 +268,12 @@ public final class AndroidTestUtil {
                   .setCodecs("avc1.64001F")
                   .build())
           .setVideoDurationUs(1_024_000L)
-          .setVideoFrameCount(30)
+          .setVideoTimestampsUs(
+              ImmutableList.of(
+                  0L, 33_366L, 66_733L, 100_100L, 133_466L, 166_833L, 200_200L, 233_566L, 266_933L,
+                  300_300L, 333_666L, 367_033L, 400_400L, 433_766L, 467_133L, 500_500L, 533_866L,
+                  567_233L, 600_600L, 633_966L, 667_333L, 700_700L, 734_066L, 767_433L, 800_800L,
+                  834_166L, 867_533L, 900_900L, 934_266L, 967_633L))
           .build();
 
   public static final AssetInfo BT601_MOV_ASSET =

libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlaybackTest.java

@@ -21,27 +21,36 @@ import static androidx.media3.test.utils.BitmapPixelTestUtil.createArgb8888Bitma
 import static androidx.media3.test.utils.BitmapPixelTestUtil.getBitmapAveragePixelAbsoluteDifferenceArgb8888;
 import static androidx.media3.test.utils.BitmapPixelTestUtil.readBitmap;
 import static androidx.media3.transformer.AndroidTestUtil.MP4_ASSET;
+import static androidx.media3.transformer.AndroidTestUtil.PNG_ASSET;
 import static androidx.media3.transformer.mh.performance.PlaybackTestUtil.createTimestampOverlay;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import android.app.Instrumentation;
+import android.content.Context;
 import android.graphics.Bitmap;
 import android.graphics.PixelFormat;
 import android.media.Image;
 import android.media.ImageReader;
+import androidx.media3.common.Effect;
 import androidx.media3.common.MediaItem;
 import androidx.media3.common.util.ConditionVariable;
 import androidx.media3.common.util.Size;
 import androidx.media3.common.util.Util;
+import androidx.media3.effect.GlEffect;
 import androidx.media3.transformer.Composition;
 import androidx.media3.transformer.CompositionPlayer;
 import androidx.media3.transformer.EditedMediaItem;
 import androidx.media3.transformer.EditedMediaItemSequence;
 import androidx.media3.transformer.Effects;
+import androidx.media3.transformer.InputTimestampRecordingShaderProgram;
+import androidx.media3.transformer.PlayerTestListener;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.platform.app.InstrumentationRegistry;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
@@ -59,9 +68,21 @@ public class CompositionPlaybackTest {
   private static final String TEST_DIRECTORY = "test-generated-goldens/ExoPlayerPlaybackTest";
   private static final long TEST_TIMEOUT_MS = 10_000;
 
+  private static final MediaItem VIDEO_MEDIA_ITEM = MediaItem.fromUri(MP4_ASSET.uri);
+  private static final long VIDEO_DURATION_US = MP4_ASSET.videoDurationUs;
+  private static final ImmutableList<Long> VIDEO_TIMESTAMPS_US = MP4_ASSET.videoTimestampsUs;
+  private static final MediaItem IMAGE_MEDIA_ITEM =
+      new MediaItem.Builder().setUri(PNG_ASSET.uri).setImageDurationMs(200).build();
+  private static final long IMAGE_DURATION_US = 200_000;
+  // 200 ms at 30 fps (default frame rate)
+  private static final ImmutableList<Long> IMAGE_TIMESTAMPS_US =
+      ImmutableList.of(0L, 33_333L, 66_667L, 100_000L, 133_333L, 166_667L);
+
   @Rule public final TestName testName = new TestName();
 
-  private final Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
+  private final Context context = getInstrumentation().getContext().getApplicationContext();
+  private final PlayerTestListener playerTestListener = new PlayerTestListener(TEST_TIMEOUT_MS);
+
   private @MonotonicNonNull CompositionPlayer player;
   private @MonotonicNonNull ImageReader outputImageReader;
   private String testId;
@@ -73,7 +94,8 @@ public class CompositionPlaybackTest {
 
   @After
   public void tearDown() {
-    instrumentation.runOnMainSync(
+    getInstrumentation()
+        .runOnMainSync(
             () -> {
               if (player != null) {
                 player.release();
@@ -95,9 +117,10 @@ public class CompositionPlaybackTest {
             PixelFormat.RGBA_8888,
             /* maxImages= */ 1);
 
-    instrumentation.runOnMainSync(
+    getInstrumentation()
+        .runOnMainSync(
             () -> {
-          player = new CompositionPlayer.Builder(instrumentation.getContext()).build();
+              player = new CompositionPlayer.Builder(context).build();
               outputImageReader.setOnImageAvailableListener(
                   imageReader -> {
                     try (Image image = imageReader.acquireLatestImage()) {
@@ -141,4 +164,178 @@ public class CompositionPlaybackTest {
     assertThat(averagePixelAbsoluteDifference).isAtMost(MAXIMUM_AVERAGE_PIXEL_ABSOLUTE_DIFFERENCE);
     // TODO: b/315800590 - Verify onFirstFrameRendered is invoked only once.
   }
+
+  @Test
+  public void playback_sequenceOfVideos_effectsReceiveCorrectTimestamps() throws Exception {
+    InputTimestampRecordingShaderProgram inputTimestampRecordingShaderProgram =
+        new InputTimestampRecordingShaderProgram();
+    Effect videoEffect = (GlEffect) (context, useHdr) -> inputTimestampRecordingShaderProgram;
+    EditedMediaItem editedMediaItem =
+        new EditedMediaItem.Builder(VIDEO_MEDIA_ITEM)
+            .setDurationUs(VIDEO_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    Composition composition =
+        new Composition.Builder(
+                new EditedMediaItemSequence(ImmutableList.of(editedMediaItem, editedMediaItem)))
+            .build();
+    List<Long> expectedTimestampsUs = new ArrayList<>();
+    expectedTimestampsUs.addAll(VIDEO_TIMESTAMPS_US);
+    expectedTimestampsUs.addAll(
+        Lists.newArrayList(
+            Iterables.transform(
+                VIDEO_TIMESTAMPS_US, timestampUs -> (VIDEO_DURATION_US + timestampUs))));
+
+    getInstrumentation()
+        .runOnMainSync(
+            () -> {
+              player = new CompositionPlayer.Builder(context).build();
+              player.addListener(playerTestListener);
+              player.setComposition(composition);
+              player.prepare();
+              player.play();
+            });
+    playerTestListener.waitUntilPlayerEnded();
+
+    assertThat(inputTimestampRecordingShaderProgram.getInputTimestampsUs())
+        .isEqualTo(expectedTimestampsUs);
+  }
+
+  @Test
+  public void playback_sequenceOfImages_effectsReceiveCorrectTimestamps() throws Exception {
+    InputTimestampRecordingShaderProgram inputTimestampRecordingShaderProgram =
+        new InputTimestampRecordingShaderProgram();
+    Effect videoEffect = (GlEffect) (context, useHdr) -> inputTimestampRecordingShaderProgram;
+    EditedMediaItem editedMediaItem =
+        new EditedMediaItem.Builder(IMAGE_MEDIA_ITEM)
+            .setDurationUs(IMAGE_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    Composition composition =
+        new Composition.Builder(
+                new EditedMediaItemSequence(ImmutableList.of(editedMediaItem, editedMediaItem)))
+            .build();
+    List<Long> expectedTimestampsUs = new ArrayList<>();
+    expectedTimestampsUs.addAll(IMAGE_TIMESTAMPS_US);
+    expectedTimestampsUs.addAll(
+        Lists.newArrayList(
+            Iterables.transform(
+                IMAGE_TIMESTAMPS_US, timestampUs -> (IMAGE_DURATION_US + timestampUs))));
+
+    getInstrumentation()
+        .runOnMainSync(
+            () -> {
+              player = new CompositionPlayer.Builder(context).build();
+              player.addListener(playerTestListener);
+              player.setComposition(composition);
+              player.prepare();
+              player.play();
+            });
+    playerTestListener.waitUntilPlayerEnded();
+
+    assertThat(inputTimestampRecordingShaderProgram.getInputTimestampsUs())
+        .isEqualTo(expectedTimestampsUs);
+  }
+
+  @Test
+  public void playback_sequenceOfVideoAndImage_effectsReceiveCorrectTimestamps() throws Exception {
+    InputTimestampRecordingShaderProgram inputTimestampRecordingShaderProgram =
+        new InputTimestampRecordingShaderProgram();
+    Effect videoEffect = (GlEffect) (context, useHdr) -> inputTimestampRecordingShaderProgram;
+    EditedMediaItem videoEditedMediaItem =
+        new EditedMediaItem.Builder(VIDEO_MEDIA_ITEM)
+            .setDurationUs(VIDEO_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    EditedMediaItem imageEditedMediaItem =
+        new EditedMediaItem.Builder(IMAGE_MEDIA_ITEM)
+            .setDurationUs(IMAGE_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    Composition composition =
+        new Composition.Builder(
+                new EditedMediaItemSequence(
+                    ImmutableList.of(videoEditedMediaItem, imageEditedMediaItem)))
+            .build();
+    List<Long> expectedTimestampsUs = new ArrayList<>();
+    expectedTimestampsUs.addAll(VIDEO_TIMESTAMPS_US);
+    expectedTimestampsUs.addAll(
+        Lists.newArrayList(
+            Iterables.transform(
+                IMAGE_TIMESTAMPS_US, timestampUs -> (VIDEO_DURATION_US + timestampUs))));
+
+    getInstrumentation()
+        .runOnMainSync(
+            () -> {
+              player = new CompositionPlayer.Builder(context).build();
+              player.addListener(playerTestListener);
+              player.setComposition(composition);
+              player.prepare();
+              player.play();
+            });
+    playerTestListener.waitUntilPlayerEnded();
+
+    assertThat(inputTimestampRecordingShaderProgram.getInputTimestampsUs())
+        .isEqualTo(expectedTimestampsUs);
+  }
+
+  @Test
+  public void playback_sequenceOfImageAndVideo_effectsReceiveCorrectTimestamps() throws Exception {
+    InputTimestampRecordingShaderProgram inputTimestampRecordingShaderProgram =
+        new InputTimestampRecordingShaderProgram();
+    Effect videoEffect = (GlEffect) (context, useHdr) -> inputTimestampRecordingShaderProgram;
+    EditedMediaItem imageEditedMediaItem =
+        new EditedMediaItem.Builder(IMAGE_MEDIA_ITEM)
+            .setDurationUs(IMAGE_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    EditedMediaItem videoEditedMediaItem =
+        new EditedMediaItem.Builder(VIDEO_MEDIA_ITEM)
+            .setDurationUs(VIDEO_DURATION_US)
+            .setEffects(
+                new Effects(
+                    /* audioProcessors= */ ImmutableList.of(),
+                    /* videoEffects= */ ImmutableList.of(videoEffect)))
+            .build();
+    Composition composition =
+        new Composition.Builder(
+                new EditedMediaItemSequence(
+                    ImmutableList.of(imageEditedMediaItem, videoEditedMediaItem)))
+            .build();
+    List<Long> expectedTimestampsUs = new ArrayList<>();
+    expectedTimestampsUs.addAll(IMAGE_TIMESTAMPS_US);
+    expectedTimestampsUs.addAll(
+        Lists.newArrayList(
+            Iterables.transform(
+                VIDEO_TIMESTAMPS_US, timestampUs -> (IMAGE_DURATION_US + timestampUs))));
+
+    getInstrumentation()
+        .runOnMainSync(
+            () -> {
+              player = new CompositionPlayer.Builder(context).build();
+              player.addListener(playerTestListener);
+              player.setComposition(composition);
+              player.prepare();
+              player.play();
+            });
+    playerTestListener.waitUntilPlayerEnded();
+
+    assertThat(inputTimestampRecordingShaderProgram.getInputTimestampsUs())
+        .isEqualTo(expectedTimestampsUs);
+  }
 }

libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlayerSeekTest.java

@@ -64,15 +64,11 @@ public class CompositionPlayerSeekTest {
 
   private static final MediaItem VIDEO_MEDIA_ITEM = MediaItem.fromUri(MP4_ASSET.uri);
   private static final long VIDEO_DURATION_US = MP4_ASSET.videoDurationUs;
-  private static final ImmutableList<Long> VIDEO_TIMESTAMPS_US =
+  private static final ImmutableList<Long> VIDEO_TIMESTAMPS_US = MP4_ASSET.videoTimestampsUs;
-      ImmutableList.of(
-          0L, 33_366L, 66_733L, 100_100L, 133_466L, 166_833L, 200_200L, 233_566L, 266_933L,
-          300_300L, 333_666L, 367_033L, 400_400L, 433_766L, 467_133L, 500_500L, 533_866L, 567_233L,
-          600_600L, 633_966L, 667_333L, 700_700L, 734_066L, 767_433L, 800_800L, 834_166L, 867_533L,
-          900_900L, 934_266L, 967_633L);
   private static final MediaItem IMAGE_MEDIA_ITEM =
       new MediaItem.Builder().setUri(PNG_ASSET.uri).setImageDurationMs(200).build();
   private static final long IMAGE_DURATION_US = 200_000;
+  // 200 ms at 30 fps (default frame rate)
   private static final ImmutableList<Long> IMAGE_TIMESTAMPS_US =
       ImmutableList.of(0L, 33_333L, 66_667L, 100_000L, 133_333L, 166_667L);
 
@@ -80,16 +76,16 @@ public class CompositionPlayerSeekTest {
   public ActivityScenarioRule<SurfaceTestActivity> rule =
       new ActivityScenarioRule<>(SurfaceTestActivity.class);
 
-  private Context applicationContext;
+  private final Context applicationContext =
-  private PlayerTestListener playerTestListener;
+      getInstrumentation().getContext().getApplicationContext();
+  private final PlayerTestListener playerTestListener = new PlayerTestListener(TEST_TIMEOUT_MS);
+
   private CompositionPlayer compositionPlayer;
   private SurfaceView surfaceView;
 
   @Before
   public void setUp() {
     rule.getScenario().onActivity(activity -> surfaceView = activity.getSurfaceView());
-    applicationContext = getInstrumentation().getContext().getApplicationContext();
-    playerTestListener = new PlayerTestListener(TEST_TIMEOUT_MS);
   }
 
   @After