Implement retry flow on DriveUploader.

chrome/browser/google_apis/drive_uploader_unittest.cc

Index: chrome/browser/google_apis/drive_uploader_unittest.cc
diff --git a/chrome/browser/google_apis/drive_uploader_unittest.cc b/chrome/browser/google_apis/drive_uploader_unittest.cc
index 38aafedabe94ade969d7f689a553580cbac89fcc..2209a7dc03a18cfb16a46b0f8b03737cbf9ffc46 100644
--- a/chrome/browser/google_apis/drive_uploader_unittest.cc
+++ b/chrome/browser/google_apis/drive_uploader_unittest.cc
@@ -58,12 +58,22 @@ class MockDriveServiceWithUploadExpectation : public DummyDriveService {
       const std::string& expected_upload_content)
      : expected_upload_content_(expected_upload_content),
        received_bytes_(0),
-       resume_upload_call_count_(0) {}
+       resume_upload_call_count_(0),
+       get_upload_status_call_count_(0),
+       resume_upload_failure_point_(-1) {}
 
   int64 received_bytes() const { return received_bytes_; }
 
   int64 resume_upload_call_count() const { return resume_upload_call_count_; }
 
+  int64 get_upload_status_call_count() const {
+    return get_upload_status_call_count_;
+  }
+
+  void set_resume_upload_failure_point(int64 resume_upload_failure_point) {
+    resume_upload_failure_point_ = resume_upload_failure_point;
+  }
+
  private:
   // DriveServiceInterface overrides.
   // Handles a request for obtaining an upload location URL.
@@ -122,6 +132,17 @@ class MockDriveServiceWithUploadExpectation : public DummyDriveService {
 
     // Update the internal status of the current upload session.
     resume_upload_call_count_++;
+    if (resume_upload_call_count_ == resume_upload_failure_point_) {
+      // Emulate the server error.
+      // ResumeUpload is an asynchronous function, so don't callback directly.
+      scoped_ptr<ResourceEntry> entry;
+      MessageLoop::current()->PostTask(
+          FROM_HERE,
+          base::Bind(callback,
+                     UploadRangeResponse(HTTP_SERVICE_UNAVAILABLE, 0, 0),
+                     base::Passed(&entry)));
+      return;
+    }
     received_bytes_ = params.end_position;
 
     // Callback with response.
@@ -144,9 +165,50 @@ class MockDriveServiceWithUploadExpectation : public DummyDriveService {
         base::Bind(callback, response, base::Passed(&entry)));
   }
 
+  // Handles a request for uploading a chunk of bytes.
+  virtual void GetUploadStatus(
+      UploadMode upload_mode,
+      const FilePath& drive_file_path,
+      const GURL& upload_url,
+      int64 content_length,
+      const UploadRangeCallback& callback) OVERRIDE {
+    EXPECT_EQ(GURL(kTestUploadURL), upload_url);
+    const int64 expected_upload_content_size = expected_upload_content_.size();
+    EXPECT_EQ(expected_upload_content_size, content_length);
+
+    // Update the internal status of the current upload session.
+    get_upload_status_call_count_++;
+
+    // Callback with response.
+    UploadRangeResponse response;
+    scoped_ptr<ResourceEntry> entry;
+    if (received_bytes_ == content_length) {
+      response = UploadRangeResponse(
+          upload_mode == UPLOAD_NEW_FILE ? HTTP_CREATED : HTTP_SUCCESS,
+          -1, -1);
+
+      base::DictionaryValue dict;
+      dict.Set("id.$t", new base::StringValue(kTestDummyId));
+      entry = ResourceEntry::CreateFrom(dict);
+    } else {
+      response =
+          UploadRangeResponse(HTTP_RESUME_INCOMPLETE, 0, received_bytes_);
+    }
+    // GetUploadStatus is an asynchronous function, so don't callback directly.
+    MessageLoop::current()->PostTask(FROM_HERE,
+        base::Bind(callback, response, base::Passed(&entry)));
+  }
+
   std::string expected_upload_content_;
   int64 received_bytes_;
   int64 resume_upload_call_count_;
+  int64 get_upload_status_call_count_;
+
+  // When the resume_upload_call_count is equal to this value,
+  // this mock server will return 503 SERVICE_UNAVAILABLE to test the
+  // failure case. To disable this function, set the value to -1 so that
+  // the resume_upload_call_count_ won't be equal to this in real use cases.
+  int64 resume_upload_failure_point_;
 };
 
 // Mock DriveService that returns a failure at InitiateUpload().
@@ -344,6 +406,37 @@ TEST_F(DriveUploaderTest, UploadNew1234KB) {
   EXPECT_EQ(kTestDummyId, out.resource_entry->id());
 }
 
+TEST_F(DriveUploaderTest, UploadNew1234KBWithRetry) {
+  FilePath local_path;
+  std::string data;
+  ASSERT_TRUE(CreateFileOfSpecifiedSize(temp_dir_.path(), 1234 * 1024,
+                                        &local_path, &data));
+  UploadCompletionCallbackResult out;
+
+  MockDriveServiceWithUploadExpectation mock_service(data);
+  // Emulate the server failure when ResumeUpload is invoked twice.
+  mock_service.set_resume_upload_failure_point(2);
+  DriveUploader uploader(&mock_service);
+  uploader.UploadNewFile(
+      GURL(kTestInitialUploadURL),
+      FilePath::FromUTF8Unsafe(kTestDrivePath),
+      local_path,
+      kTestDocumentTitle,
+      kTestMimeType,
+      base::Bind(&CopyResultsFromUploadCompletionCallbackAndQuit, &out));
+  message_loop_.Run();
+
+  // The file should be split into 3 chunks (1234 = 512 + 512 + 210) + 1 retry.
+  EXPECT_EQ(4, mock_service.resume_upload_call_count());
+  EXPECT_EQ(1, mock_service.get_upload_status_call_count());
+  EXPECT_EQ(1234 * 1024, mock_service.received_bytes());
+  EXPECT_EQ(DRIVE_UPLOAD_OK, out.error);
+  EXPECT_EQ(FilePath::FromUTF8Unsafe(kTestDrivePath), out.drive_path);
+  EXPECT_EQ(local_path, out.file_path);
+  ASSERT_TRUE(out.resource_entry);
+  EXPECT_EQ(kTestDummyId, out.resource_entry->id());
+}
+
 TEST_F(DriveUploaderTest, InitiateUploadFail) {
   FilePath local_path;
   std::string data;