From 4ba9f59492f5d2c97caa9981e3073763b204b861 Mon Sep 17 00:00:00 2001
From: kimvde <kimvde@google.com>
Date: Thu, 22 Aug 2024 00:12:53 -0700
Subject: [PATCH] Improve logic to wait for a seek in CompositionPlayerSeekTest

The logic was assuming that the shader program was only flushed when
seeking. This is true if a single renderer is used but, with multiple
renderers, the shader program can be flushed at the transition.

This change is necessary to make the test pass for prewarming because 2
video renderers will be used in that case.

PiperOrigin-RevId: 666215967
---
 .../mh/performance/CompositionPlayerSeekTest.java | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlayerSeekTest.java b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlayerSeekTest.java
index 75eef0ab52..1e1ada4043 100644
--- a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlayerSeekTest.java
+++ b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/performance/CompositionPlayerSeekTest.java
@@ -133,8 +133,7 @@ public class CompositionPlayerSeekTest {
             () -> {
               compositionPlayer = new CompositionPlayer.Builder(applicationContext).build();
               // Set a surface on the player even though there is no UI on this test. We need a
-              // surface
-              // otherwise the player will skip/drop video frames.
+              // surface otherwise the player will skip/drop video frames.
               compositionPlayer.setVideoSurfaceView(surfaceView);
               compositionPlayer.addListener(playerTestListener);
               compositionPlayer.setComposition(
@@ -429,8 +428,13 @@ public class CompositionPlayerSeekTest {
           @Override
           public void flush() {
             super.flush();
-            shaderProgramShouldBlockInput.set(false);
-            framesReceivedLatch.reset(Integer.MAX_VALUE);
+            if (framesReceivedLatch.getCount() == 0) {
+              // The flush is caused by the seek operation. We do this check because the shader
+              // program can be flushed for other reasons, for example at the transition between 2
+              // renderers.
+              shaderProgramShouldBlockInput.set(false);
+              framesReceivedLatch.reset(Integer.MAX_VALUE);
+            }
           }
         };
 
@@ -449,8 +453,7 @@ public class CompositionPlayerSeekTest {
             () -> {
               compositionPlayer = new CompositionPlayer.Builder(applicationContext).build();
               // Set a surface on the player even though there is no UI on this test. We need a
-              // surface
-              // otherwise the player will skip/drop video frames.
+              // surface otherwise the player will skip/drop video frames.
               compositionPlayer.setVideoSurfaceView(surfaceView);
               compositionPlayer.addListener(playerTestListener);
               compositionPlayer.setComposition(