Record multiple video parts for UI tests.

Bug: 144698400
Test: included
Change-Id: I0f5880867ff89e36e6fddbf0d2423521837e0dff
diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/ScreenRecordCollector.java b/libraries/device-collectors/src/main/java/android/device/collectors/ScreenRecordCollector.java
index 86abe33..69f7085 100644
--- a/libraries/device-collectors/src/main/java/android/device/collectors/ScreenRecordCollector.java
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/ScreenRecordCollector.java
@@ -37,6 +37,7 @@
  */
 @OptionClass(alias = "screen-record-collector")
 public class ScreenRecordCollector extends BaseMetricListener {
+    @VisibleForTesting static final int MAX_RECORDING_PARTS = 5;
     private static final long VIDEO_TAIL_BUFFER = 2000;
 
     static final String OUTPUT_DIR = "run_listeners/videos";
@@ -44,6 +45,11 @@
     private UiDevice mDevice;
     private File mDestDir;
 
+    // Tracks multiple parts to a single recording.
+    private int mParts;
+    // Avoid recording after the test is finished.
+    private boolean mContinue;
+
     // Tracks the test iterations to ensure that each failure gets unique filenames.
     // Key: test description; value: number of iterations.
     private Map<String, Integer> mTestIterations = new HashMap<String, Integer>();
@@ -62,7 +68,9 @@
         // Track the number of iteration for this test.
         amendIterations(description);
         // Start the screen recording operation.
-        startScreenRecordThread(getOutputFile(description).getAbsolutePath());
+        mParts = 1;
+        mContinue = true;
+        startScreenRecordThread(description);
     }
 
     @Override
@@ -75,11 +83,14 @@
         // Add some extra time to the video end.
         SystemClock.sleep(getTailBuffer());
         // Ctrl + C all screen record processes.
+        mContinue = false;
         killScreenRecordProcesses();
 
-        // Add the output file to the data record.
-        File output = getOutputFile(description);
-        testData.addFileMetric(String.format("%s_%s", getTag(), output.getName()), output);
+        // Add the output files to the data record.
+        for (int i = 1; i < mParts; i++) {
+            File output = getOutputFile(description, i);
+            testData.addFileMetric(String.format("%s_%s", getTag(), output.getName()), output);
+        }
 
         // TODO(b/144869954): Delete when tests pass.
     }
@@ -91,31 +102,36 @@
         mTestIterations.computeIfAbsent(testName, name -> 1);
     }
 
-    private File getOutputFile(Description description) {
+    private File getOutputFile(Description description, int part) {
         final String baseName =
                 String.format("%s.%s", description.getClassName(), description.getMethodName());
         // Omit the iteration number for the first iteration.
         int iteration = mTestIterations.get(description.getDisplayName());
         final String fileName =
                 String.format(
-                        "%s-video.mp4",
+                        "%s-video%s.mp4",
                         iteration == 1
                                 ? baseName
-                                : String.join("-", baseName, String.valueOf(iteration)));
+                                : String.join("-", baseName, String.valueOf(iteration)),
+                        part == 1 ? "" : part);
         return Paths.get(mDestDir.getAbsolutePath(), fileName).toFile();
     }
 
     /** Spawns a thread to start screen recording that will save to the provided {@code path}. */
-    public void startScreenRecordThread(String path) {
-        Log.d(getTag(), String.format("Recording screen to %s", path));
+    public void startScreenRecordThread(final Description description) {
         new Thread("test-screenrecord-thread") {
             @Override
             public void run() {
                 try {
-                    // Make sure not to block on this background command.
-                    getDevice().executeShellCommand(String.format("screenrecord %s", path));
+                    for (int i = 0; i < MAX_RECORDING_PARTS && mContinue; i++) {
+                        String output = getOutputFile(description, mParts).getAbsolutePath();
+                        Log.d(getTag(), String.format("Recording screen to %s", output));
+                        // Make sure not to block on this background command so the test runs.
+                        getDevice().executeShellCommand(String.format("screenrecord %s", output));
+                        mParts++;
+                    }
                 } catch (IOException e) {
-                    throw new RuntimeException("Failed to start screen recording.");
+                    throw new RuntimeException("Caught exception while screen recording.");
                 }
             }
         }.start();
diff --git a/libraries/device-collectors/src/test/java/android/device/collectors/ScreenRecordCollectorTest.java b/libraries/device-collectors/src/test/java/android/device/collectors/ScreenRecordCollectorTest.java
index d7198c1..8f45a74 100644
--- a/libraries/device-collectors/src/test/java/android/device/collectors/ScreenRecordCollectorTest.java
+++ b/libraries/device-collectors/src/test/java/android/device/collectors/ScreenRecordCollectorTest.java
@@ -17,6 +17,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.AdditionalMatchers.not;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.endsWith;
 import static org.mockito.ArgumentMatchers.eq;
@@ -112,10 +113,16 @@
         for (int i = 1; i <= NUM_TEST_CASE; i++) {
             // Verify a thread is started when the test starts.
             mListener.testStarted(mTestDesc);
-            verify(mListener, times(i)).startScreenRecordThread(anyString());
+            verify(mListener, times(i)).startScreenRecordThread(any());
             // Delay verification by 100 ms to ensure the thread was started.
             SystemClock.sleep(100);
-            verify(mDevice, times(i)).executeShellCommand(matches("screenrecord .*"));
+            // Expect all recordings to be finished because of mocked commands.
+            verify(mDevice, times(i)).executeShellCommand(matches("screenrecord .*video.mp4"));
+            for (int r = 2; r < ScreenRecordCollector.MAX_RECORDING_PARTS; r++) {
+                verify(mDevice, times(i))
+                        .executeShellCommand(
+                                matches(String.format("screenrecord .*video%d.mp4", r)));
+            }
 
             // Alternate between pass and fail for variety.
             if (i % 2 == 0) {
@@ -149,34 +156,42 @@
                 if (key.contains("mp4")) videoCount++;
             }
         }
-        assertEquals(NUM_TEST_CASE, videoCount);
+        assertEquals(NUM_TEST_CASE * ScreenRecordCollector.MAX_RECORDING_PARTS, videoCount);
     }
 
     /** Test that screen recording is properly done for multiple tests and labels iterations. */
     @Test
-    public void testMultipleScreenRecords() throws Exception {
+    public void testScreenRecord_multipleTests() throws Exception {
         mListener = initListener();
 
         // Run through a sequence of `NUM_TEST_CASE` failing tests.
         mListener.testRunStarted(mRunDesc);
-        verify(mListener).createAndEmptyDirectory(ScreenRecordCollector.OUTPUT_DIR);
 
         // Walk through a number of test cases to simulate behavior.
         for (int i = 1; i <= NUM_TEST_CASE; i++) {
             mListener.testStarted(mTestDesc);
+            SystemClock.sleep(100);
             mListener.testFinished(mTestDesc);
         }
         mListener.testRunFinished(new Result());
 
         // Verify that videos are saved with iterations.
-        InOrder videoVerifier = inOrder(mListener);
+        InOrder videoVerifier = inOrder(mDevice);
         // The first video should not have an iteration number.
-        videoVerifier.verify(mListener).startScreenRecordThread(matches("^.*[^1].mp4$"));
+        videoVerifier
+                .verify(mDevice, times(ScreenRecordCollector.MAX_RECORDING_PARTS))
+                .executeShellCommand(matches("^.*[^1]-video.*.mp4$"));
         // The subsequent videos should have an iteration number.
         for (int i = 1; i < NUM_TEST_CASE; i++) {
             videoVerifier
-                    .verify(mListener)
-                    .startScreenRecordThread(endsWith(String.format("%d-video.mp4", i + 1)));
+                    .verify(mDevice)
+                    .executeShellCommand(endsWith(String.format("%d-video.mp4", i + 1)));
+            // Verify the iteration-specific and part-specific interactions too.
+            for (int p = 2; p <= ScreenRecordCollector.MAX_RECORDING_PARTS; p++) {
+                videoVerifier
+                        .verify(mDevice)
+                        .executeShellCommand(endsWith(String.format("%d-video%d.mp4", i + 1, p)));
+            }
         }
     }
 }