Add a unit test to check partition info

Change-Id: I397d7005961a037c9c9cb29e3ff0a3d39a501d15

1 files changed, 52 insertions, 0 deletions

diff --git a/test/simple_encode_test.cc b/test/simple_encode_test.cc
index d917c2403..e4bb1ebcf 100644
--- a/test/simple_encode_test.cc
+++ b/test/simple_encode_test.cc
@@ -235,6 +235,58 @@ TEST(SimpleEncode, EncodeConsistencyTest2) {
   simple_encode_2.EndEncode();
 }
 
+// Test the information stored in encoder is the same between two encode runs.
+TEST(SimpleEncode, EncodeConsistencyTest3) {
+  std::vector<int> quantize_index_list;
+  const int num_rows_4x4 = GetNumUnit4x4(w);
+  const int num_cols_4x4 = GetNumUnit4x4(h);
+  const int num_units_4x4 = num_rows_4x4 * num_cols_4x4;
+  // The first encode.
+  SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
+                             target_bitrate, num_frames, infile_path);
+  simple_encode.ComputeFirstPassStats();
+  const int num_coding_frames = simple_encode.GetCodingFrameNum();
+  std::vector<PartitionInfo> partition_info_list(num_units_4x4 *
+                                                 num_coding_frames);
+  simple_encode.StartEncode();
+  for (int i = 0; i < num_coding_frames; ++i) {
+    EncodeFrameResult encode_frame_result;
+    simple_encode.EncodeFrame(&encode_frame_result);
+    quantize_index_list.push_back(encode_frame_result.quantize_index);
+    for (int j = 0; j < num_rows_4x4 * num_cols_4x4; ++j) {
+      partition_info_list[i * num_units_4x4 + j] =
+          encode_frame_result.partition_info[j];
+    }
+  }
+  simple_encode.EndEncode();
+  // The second encode.
+  SimpleEncode simple_encode_2(w, h, frame_rate_num, frame_rate_den,
+                               target_bitrate, num_frames, infile_path);
+  simple_encode_2.ComputeFirstPassStats();
+  const int num_coding_frames_2 = simple_encode_2.GetCodingFrameNum();
+  simple_encode_2.StartEncode();
+  for (int i = 0; i < num_coding_frames_2; ++i) {
+    EncodeFrameResult encode_frame_result;
+    simple_encode_2.EncodeFrameWithQuantizeIndex(&encode_frame_result,
+                                                 quantize_index_list[i]);
+    for (int j = 0; j < num_rows_4x4 * num_cols_4x4; ++j) {
+      EXPECT_EQ(encode_frame_result.partition_info[j].row,
+                partition_info_list[i * num_units_4x4 + j].row);
+      EXPECT_EQ(encode_frame_result.partition_info[j].column,
+                partition_info_list[i * num_units_4x4 + j].column);
+      EXPECT_EQ(encode_frame_result.partition_info[j].row_start,
+                partition_info_list[i * num_units_4x4 + j].row_start);
+      EXPECT_EQ(encode_frame_result.partition_info[j].column_start,
+                partition_info_list[i * num_units_4x4 + j].column_start);
+      EXPECT_EQ(encode_frame_result.partition_info[j].width,
+                partition_info_list[i * num_units_4x4 + j].width);
+      EXPECT_EQ(encode_frame_result.partition_info[j].height,
+                partition_info_list[i * num_units_4x4 + j].height);
+    }
+  }
+  simple_encode_2.EndEncode();
+}
+
 TEST(SimpleEncode, GetEncodeFrameInfo) {
   // Makes sure that the encode_frame_info obtained from GetEncodeFrameInfo()
   // matches the counterpart in encode_frame_result obtained from EncodeFrame()