Implement retry flow on DriveUploader.

chrome/browser/google_apis/drive_uploader_unittest.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome/browser/google_apis/drive_uploader.h"
 
 #include <algorithm>
 #include <cstdlib>
 #include <string>
 
@@ skipping 40 matching lines @@
 // Mock DriveService that verifies if the uploaded content matches the preset
 // expectation.
 class MockDriveServiceWithUploadExpectation : public DummyDriveService {
  public:
   // Sets up an expected upload content. InitiateUpload and ResumeUpload will
   // verify that the specified data is correctly uploaded.
   explicit MockDriveServiceWithUploadExpectation(
       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.
   virtual void InitiateUpload(const InitiateUploadParams& params,
                               const InitiateUploadCallback& callback) OVERRIDE {
     const int64 expected_size = expected_upload_content_.size();
 
     // Verify that the expected parameters are passed.
     if (params.upload_mode == UPLOAD_NEW_FILE)
       EXPECT_EQ(kTestDocumentTitle, params.title);
@@ skipping 38 matching lines @@
 
     // The upload URL returned by InitiateUpload() must be used.
     EXPECT_EQ(GURL(kTestUploadURL), params.upload_location);
 
     // Other parameters should be the exact values passed to DriveUploader.
     EXPECT_EQ(expected_size, params.content_length);
     EXPECT_EQ(kTestMimeType, params.content_type);
 
     // 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.
     UploadRangeResponse response;
     scoped_ptr<ResourceEntry> entry;
     if (received_bytes_ == params.content_length) {
       response = UploadRangeResponse(
           params.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,
                                       params.end_position);
     }
     // ResumeUpload is an asynchronous function, so don't callback directly.
     MessageLoop::current()->PostTask(FROM_HERE,
         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().
 class MockDriveServiceNoConnectionAtInitiate : public DummyDriveService {
   // Returns error.
   virtual void InitiateUpload(const InitiateUploadParams& params,
                               const InitiateUploadCallback& callback) OVERRIDE {
     MessageLoop::current()->PostTask(FROM_HERE,
         base::Bind(callback, GDATA_NO_CONNECTION, GURL()));
   }
@@ skipping 177 matching lines @@
   // The file should be split into 3 chunks (1234 = 512 + 512 + 210).
   EXPECT_EQ(3, mock_service.resume_upload_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, 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;
   ASSERT_TRUE(CreateFileOfSpecifiedSize(temp_dir_.path(), 512 * 1024,
                                         &local_path, &data));
 
   UploadCompletionCallbackResult out;
 
   MockDriveServiceNoConnectionAtInitiate mock_service;
   DriveUploader uploader(&mock_service);
@@ skipping 84 matching lines @@
       kTestMimeType,
       "",  // etag
       base::Bind(&CopyResultsFromUploadCompletionCallbackAndQuit, &out));
   message_loop_.Run();
 
   // Should return failure without doing any attempt to connect to the server.
   EXPECT_EQ(DRIVE_UPLOAD_ERROR_NOT_FOUND, out.error);
 }
 
 }  // namespace google_apis