Fix two-pass encoding and create basic CLI tests

.gitattributes

@@ -0,0 +1 @@
+*.y4m binary

Cargo.toml

@@ -108,6 +108,7 @@ features = ["parallel"]
 signal-hook = { version = "0.1.9", optional = true }
 
 [dev-dependencies]
+assert_cmd = "0.12"
 criterion = "0.3"
 pretty_assertions = "0.6"
 interpolate_name = "0.2.2"

src/bin/rav1e.rs

@@ -95,15 +95,14 @@ impl<D: Decoder> Source<D> {
 
 // Encode and write a frame.
 // Returns frame information in a `Result`.
-//
-// `Ok(None)` indicates that the encode is finished.
 fn process_frame<T: Pixel, D: Decoder>(
   ctx: &mut Context<T>, output_file: &mut dyn Muxer, source: &mut Source<D>,
   pass1file: Option<&mut File>, pass2file: Option<&mut File>,
   buffer: &mut [u8], buf_pos: &mut usize,
   mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>,
-) -> Result<Option<FrameSummary>, CliError> {
+) -> Result<Option<Vec<FrameSummary>>, CliError> {
   let y4m_details = source.input.get_video_details();
+  let mut frame_summaries = Vec::new();
   let mut pass1file = pass1file;
   let mut pass2file = pass2file;
   // Submit first pass data to pass 2.
@@ -143,60 +142,54 @@ fn process_frame<T: Pixel, D: Decoder>(
     }
   }
 
-  loop {
-    let pkt_wrapped = ctx.receive_packet();
-    match pkt_wrapped {
-      Ok(pkt) => {
-        output_file.write_frame(
-          pkt.input_frameno as u64,
-          pkt.data.as_ref(),
-          pkt.frame_type,
-        );
-        if let (Some(ref mut y4m_enc_uw), Some(ref rec)) =
-          (y4m_enc.as_mut(), &pkt.rec)
-        {
-          write_y4m_frame(y4m_enc_uw, rec, y4m_details);
-        }
-        return Ok(Some(pkt.into()));
-      }
-      Err(EncoderStatus::NeedMoreData) => {
-        // Read another frame then try receive_packet again
-        source.read_frame(ctx, y4m_details);
-        continue;
-      }
-      Err(EncoderStatus::EnoughData) => {
-        unreachable!();
+  let pkt_wrapped = ctx.receive_packet();
+  match pkt_wrapped {
+    Ok(pkt) => {
+      output_file.write_frame(
+        pkt.input_frameno as u64,
+        pkt.data.as_ref(),
+        pkt.frame_type,
+      );
+      if let (Some(ref mut y4m_enc_uw), Some(ref rec)) =
+        (y4m_enc.as_mut(), &pkt.rec)
+      {
+        write_y4m_frame(y4m_enc_uw, rec, y4m_details);
       }
-      Err(EncoderStatus::LimitReached) => {
-        if let Some(passfile) = pass1file.as_mut() {
-          if let Some(outbuf) = ctx.twopass_out() {
-            // The last block of data we get is the summary data that needs to go
-            //  at the start of the pass file.
-            // Seek to the start so we can write it there.
-            passfile
-              .seek(std::io::SeekFrom::Start(0))
-              .expect("Unable to seek in two-pass data file.");
-            passfile
-              .write_all(outbuf)
-              .expect("Unable to write to two-pass data file.");
-          }
+      frame_summaries.push(pkt.into());
+    }
+    Err(EncoderStatus::NeedMoreData) => {
+      source.read_frame(ctx, y4m_details);
+    }
+    Err(EncoderStatus::EnoughData) => {
+      unreachable!();
+    }
+    Err(EncoderStatus::LimitReached) => {
+      if let Some(passfile) = pass1file.as_mut() {
+        if let Some(outbuf) = ctx.twopass_out() {
+          // The last block of data we get is the summary data that needs to go
+          //  at the start of the pass file.
+          // Seek to the start so we can write it there.
+          passfile
+            .seek(std::io::SeekFrom::Start(0))
+            .expect("Unable to seek in two-pass data file.");
+          passfile
+            .write_all(outbuf)
+            .expect("Unable to write to two-pass data file.");
         }
-        // Indicate that we are finished with the encode
-        return Ok(None);
-      }
-      e @ Err(EncoderStatus::Failure) => {
-        let _ = e.map_err(|e| e.context("Failed to encode video"))?;
-      }
-      e @ Err(EncoderStatus::NotReady) => {
-        let _ = e.map_err(|e| {
-          e.context("Mismanaged handling of two-pass stats data")
-        })?;
-      }
-      Err(EncoderStatus::Encoded) => {
-        // Safely skip to the next attempt of receive_packet
       }
+      return Ok(None);
+    }
+    e @ Err(EncoderStatus::Failure) => {
+      let _ = e.map_err(|e| e.context("Failed to encode video"))?;
     }
+    e @ Err(EncoderStatus::NotReady) => {
+      let _ = e.map_err(|e| {
+        e.context("Mismanaged handling of two-pass stats data")
+      })?;
+    }
+    Err(EncoderStatus::Encoded) => {}
   }
+  Ok(Some(frame_summaries))
 }
 
 fn do_encode<T: Pixel, D: Decoder>(
@@ -233,13 +226,15 @@ fn do_encode<T: Pixel, D: Decoder>(
     y4m_enc.as_mut(),
   )? {
     if verbose != Verbose::Quiet {
-      progress.add_frame(frame_info.clone());
-      if verbose == Verbose::Verbose {
-        info!("{} - {}", frame_info, progress);
-      } else {
-        // Print a one-line progress indicator that overrides itself with every update
-        eprint!("\r{}                    ", progress);
-      };
+      for frame in frame_info {
+        progress.add_frame(frame.clone());
+        if verbose == Verbose::Verbose {
+          info!("{} - {}", frame, progress);
+        } else {
+          // Print a one-line progress indicator that overrides itself with every update
+          eprint!("\r{}                    ", progress);
+        };
+      }
 
       output.flush().unwrap();
     }

tests/.gitignore

@@ -0,0 +1 @@
+!small_input.y4m

tests/binary.rs

@@ -0,0 +1,93 @@
+#[cfg(feature = "binaries")]
+mod binary {
+  use assert_cmd::Command;
+  use rand::distributions::Alphanumeric;
+  use rand::{thread_rng, Rng};
+  use std::env::temp_dir;
+  use std::fs::File;
+  use std::io::Read;
+  use std::path::PathBuf;
+
+  fn get_y4m_input() -> Vec<u8> {
+    let mut input = File::open(&format!(
+      "{}/tests/small_input.y4m",
+      env!("CARGO_MANIFEST_DIR")
+    ))
+    .unwrap();
+    let mut data = Vec::new();
+    input.read_to_end(&mut data).unwrap();
+    data
+  }
+
+  fn get_tempfile_path(extension: &str) -> PathBuf {
+    let mut path = temp_dir();
+    let filename =
+      thread_rng().sample_iter(&Alphanumeric).take(12).collect::<String>();
+    path.push(format!("{}.{}", filename, extension));
+    path
+  }
+
+  #[test]
+  fn one_pass_qp_based() {
+    let mut cmd = Command::cargo_bin("rav1e").unwrap();
+    let outfile = get_tempfile_path("ivf");
+
+    cmd
+      .arg("--quantizer")
+      .arg("100")
+      .arg("-o")
+      .arg(&outfile)
+      .arg("-")
+      .write_stdin(get_y4m_input())
+      .assert()
+      .success();
+  }
+
+  #[test]
+  fn one_pass_bitrate_based() {
+    let mut cmd = Command::cargo_bin("rav1e").unwrap();
+    let outfile = get_tempfile_path("ivf");
+
+    cmd
+      .arg("--bitrate")
+      .arg("1000")
+      .arg("-o")
+      .arg(&outfile)
+      .arg("-")
+      .write_stdin(get_y4m_input())
+      .assert()
+      .success();
+  }
+
+  #[test]
+  fn two_pass_bitrate_based() {
+    let outfile = get_tempfile_path("ivf");
+    let passfile = get_tempfile_path("pass");
+
+    let mut cmd1 = Command::cargo_bin("rav1e").unwrap();
+    cmd1
+      .arg("--bitrate")
+      .arg("1000")
+      .arg("-o")
+      .arg(&outfile)
+      .arg("--first-pass")
+      .arg(&passfile)
+      .arg("-")
+      .write_stdin(get_y4m_input())
+      .assert()
+      .success();
+
+    let mut cmd2 = Command::cargo_bin("rav1e").unwrap();
+    cmd2
+      .arg("--bitrate")
+      .arg("1000")
+      .arg("-o")
+      .arg(&outfile)
+      .arg("--second-pass")
+      .arg(&passfile)
+      .arg("-")
+      .write_stdin(get_y4m_input())
+      .assert()
+      .success();
+  }
+}