Move the logic to determine how much data can be written to another function

content/browser/download/download_file_unittest.cc

Index: content/browser/download/download_file_unittest.cc
diff --git a/content/browser/download/download_file_unittest.cc b/content/browser/download/download_file_unittest.cc
index 055b883e9364a894b1b77c7fe48a1220d93caefe..be6fda0d620870841d7f2fba9777c9ce0baac966 100644
--- a/content/browser/download/download_file_unittest.cc
+++ b/content/browser/download/download_file_unittest.cc
@@ -396,7 +396,8 @@ class DownloadFileTest : public testing::Test {
   }
 
   // Prepare two byte streams to write to the same file sink.
-  void PrepareMultipleStreams(int64_t second_stream_length) {
+  void PrepareMultipleStreams(bool first_stream_write_all_data,
+                              int64_t second_stream_length) {
     // Create a sparse file.
     ASSERT_TRUE(CreateDownloadFile(0, true, true));
     base::FilePath initial_path(download_file_->FullPath());
@@ -405,6 +406,7 @@ class DownloadFileTest : public testing::Test {
 
     const char* stream_0_data[] = {kTestData1, kTestData2};
     const char* stream_1_data[] = {kTestData4, kTestData5};
+    const char* all_data[] = {kTestData1, kTestData2, kTestData4, kTestData5};
     size_t stream_1_offset = strlen(kTestData1) + strlen(kTestData2);
 
     // Register second SourceStream entry for the second stream.
@@ -421,8 +423,16 @@ class DownloadFileTest : public testing::Test {
 
     ::testing::Sequence s0;
     ::testing::Sequence s1;
-    SetupDataAppend(stream_1_data, 2, input_stream_1_, s1, stream_1_offset);
-    SetupDataAppend(stream_0_data, 2, input_stream_, s0, 0);
+    if (first_stream_write_all_data) {
+      SetupDataAppend(all_data, 4, input_stream_, s0, 0);
+      // The 2nd stream will abort after the first read
+      SetupDataAppend(stream_1_data, 1, input_stream_1_, s1, stream_1_offset);
+    } else {
+      SetupDataAppend(stream_0_data, 2, input_stream_, s0, 0);
+      SetupDataAppend(stream_1_data, 2, input_stream_1_, s1, stream_1_offset);
+    }
+
+
     SetupFinishStream(DOWNLOAD_INTERRUPT_REASON_NONE, input_stream_, s0);
 
     // Expectation on MockByteStreamReader for MultipleStreams tests:
@@ -433,7 +443,7 @@ class DownloadFileTest : public testing::Test {
     // The stream may terminate in the middle and less Read calls are expected.
     // 3. GetStatus: Only called if the stream is completed and last Read call
     // returns STREAM_COMPLETE.
-    if (second_stream_length == 0)
+    if (second_stream_length == 0 && !first_stream_write_all_data)
       SetupFinishStream(DOWNLOAD_INTERRUPT_REASON_NONE, input_stream_1_, s1);
     else
       EXPECT_CALL(*input_stream_1_, RegisterCallback(_)).RetiresOnSaturation();
@@ -903,7 +913,7 @@ TEST_F(DownloadFileTest, StreamNonEmptyError) {
 //
 // Activate both streams at the same time.
 TEST_F(DownloadFileTest, MutipleStreamsWrite) {
-  PrepareMultipleStreams(0);
+  PrepareMultipleStreams(false, 0);
   EXPECT_CALL(*(observer_.get()), MockDestinationCompleted(_, _));
 
   int64_t stream_0_length =
@@ -927,7 +937,7 @@ TEST_F(DownloadFileTest, MutipleStreamsWrite) {
 
 // Activate and deplete one stream, later add the second stream.
 TEST_F(DownloadFileTest, MutipleStreamsOneStreamFirst) {
-  PrepareMultipleStreams(0);
+  PrepareMultipleStreams(false, 0);
 
   int64_t stream_0_length =
       static_cast<int64_t>(strlen(kTestData1) + strlen(kTestData2));
@@ -970,7 +980,7 @@ TEST_F(DownloadFileTest, MutipleStreamsLimitedLength) {
   int64_t stream_0_length =
       static_cast<int64_t>(strlen(kTestData1) + strlen(kTestData2));
   int64_t stream_1_length = static_cast<int64_t>(strlen(kTestData4)) - 1;
-  PrepareMultipleStreams(stream_1_length);
+  PrepareMultipleStreams(false, stream_1_length);
 
   EXPECT_CALL(*(observer_.get()), MockDestinationCompleted(_, _));
 
@@ -1000,4 +1010,33 @@ TEST_F(DownloadFileTest, MutipleStreamsLimitedLength) {
   DestroyDownloadFile(0, false);
 }
 
+// Two streams write to one sink, the first stream writes the whole file before
+// the seconds stream was able to start
+TEST_F(DownloadFileTest, MutipleStreamsFirstStreamWriteAllData) {
+  PrepareMultipleStreams(true, 0);
+  int64_t stream_0_length =
+      static_cast<int64_t>(strlen(kTestData1) + strlen(kTestData2) +
+                           strlen(kTestData4) + strlen(kTestData5));
+  int64_t stream_1_length =
+      static_cast<int64_t>(strlen(kTestData4) + strlen(kTestData5));
+  sink_callback_.Run();
+  base::RunLoop().RunUntilIdle();
+
+  EXPECT_CALL(*(observer_.get()), MockDestinationCompleted(_, _));
+
+  download_file_->AddByteStream(
+      std::unique_ptr<MockByteStreamReader>(input_stream_1_),
+      stream_0_length - stream_1_length);
+  base::RunLoop().RunUntilIdle();
+
+  SourceStreamTestData stream_data_0(0, stream_0_length, true);
+  SourceStreamTestData stream_data_1(
+      stream_0_length - stream_1_length, 0, true);
+  VerifySourceStreamsStates(stream_data_0);
+  VerifySourceStreamsStates(stream_data_1);
+  EXPECT_EQ(stream_0_length, TotalBytesReceived());
+
+  DestroyDownloadFile(0);
+}
+
 }  // namespace content