Retry uploading in UploadJobImpl when error occurs

chrome/browser/chromeos/policy/upload_job_unittest.cc

 // Copyright (c) 2015 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/chromeos/policy/upload_job.h"
 
 #include <stddef.h>
 
 #include <queue>
 #include <set>
 #include <utility>
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
 #include "base/location.h"
 #include "base/macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/run_loop.h"
 #include "base/single_thread_task_runner.h"
+#include "base/test/test_simple_task_runner.h"
 #include "base/thread_task_runner_handle.h"
 #include "base/time/time.h"
 #include "chrome/browser/chromeos/policy/upload_job_impl.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "google_apis/gaia/fake_oauth2_token_service.h"
 #include "google_apis/gaia/google_service_auth_error.h"
 #include "net/http/http_status_code.h"
 #include "net/test/embedded_test_server/embedded_test_server.h"
 #include "net/test/embedded_test_server/http_request.h"
 #include "net/test/embedded_test_server/http_response.h"
 #include "net/url_request/url_request_test_util.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace policy {
 
 namespace {
 
 const char kUploadPath[] = "/upload";
 const char kRobotAccountId[] = "robot@gmail.com";
 const char kCustomField1[] = "customfield1";
 const char kCustomField2[] = "customfield2";
@@ skipping 102 matching lines @@
 void MockOAuth2TokenService::InvalidateAccessTokenImpl(
     const std::string& account_id,
     const std::string& client_id,
     const ScopeSet& scopes,
     const std::string& access_token) {
   SetTokenInvalid(access_token);
 }
 
 }  // namespace
 
+// Mock thread runner which runs each tasks immediately.
+class ImmediateThreadTaskRunner : public base::SingleThreadTaskRunner {

[Andrew T Wilson (Slow), 2016/05/10 22:00:42]

I still don't understand why this is required (why you can't use
TestSimpleTaskRunner - maybe add a comment to this effect)?

[Marton Hunyady, 2016/05/11 12:55:31]

I removed it and made it possible to set the delays for the tests separately, so
now I'm using base::ThreadTaskRunnerHandle::Get() as task runner for the tests.

The problem with TestSimpleTaskRunner is that I didn't find a way to run the
tasks posted to that automatically until run_loop_.Quit() is called from the
callbacks.

+  const scoped_refptr<base::SingleThreadTaskRunner> task_runner_;

[Andrew T Wilson (Slow), 2016/05/10 22:00:42]

What is the intended scope of this member variable? Maybe make it explicit (put
it under public or protected or whatever)

+
+ public:
+  ImmediateThreadTaskRunner()
+      : task_runner_(base::ThreadTaskRunnerHandle::Get()) {}
+
+  bool PostDelayedTask(const tracked_objects::Location& from_here,
+                       const base::Closure& task,
+                       base::TimeDelta delay) override {
+    return task_runner_->PostTask(from_here, task);
+  }
+
+  bool RunsTasksOnCurrentThread() const override {
+    return task_runner_->RunsTasksOnCurrentThread();
+  }
+
+  bool PostNonNestableDelayedTask(const tracked_objects::Location& from_here,
+                                  const base::Closure& task,
+                                  base::TimeDelta delay) override {
+    return task_runner_->PostTask(from_here, task);
+  }
+
+ protected:
+  ~ImmediateThreadTaskRunner() override {}
+
+  DISALLOW_COPY_AND_ASSIGN(ImmediateThreadTaskRunner);
+};
+
 class UploadJobTestBase : public testing::Test, public UploadJob::Delegate {
  public:
   UploadJobTestBase()
       : test_browser_thread_bundle_(
-            content::TestBrowserThreadBundle::IO_MAINLOOP) {}
+            content::TestBrowserThreadBundle::IO_MAINLOOP),
+        task_runner_(new ImmediateThreadTaskRunner) {}
 
   // policy::UploadJob::Delegate:
   void OnSuccess() override {
     if (!expected_error_)
       run_loop_.Quit();
     else
       FAIL();
   }
 
   // policy::UploadJob::Delegate:
@@ skipping 23 matching lines @@
     ASSERT_TRUE(test_server_.ShutdownAndWaitUntilComplete());
   }
 
  protected:
   std::unique_ptr<UploadJob> PrepareUploadJob(
       std::unique_ptr<UploadJobImpl::MimeBoundaryGenerator>
           mime_boundary_generator) {
     std::unique_ptr<UploadJob> upload_job(
         new UploadJobImpl(GetServerURL(), kRobotAccountId, &oauth2_service_,
                           request_context_getter_.get(), this,
-                          std::move(mime_boundary_generator)));
+                          std::move(mime_boundary_generator), task_runner_));
 
     std::map<std::string, std::string> header_entries;
     header_entries.insert(std::make_pair(kCustomField1, "CUSTOM1"));
     std::unique_ptr<std::string> data(new std::string(kTestPayload1));
     upload_job->AddDataSegment("Name1", "file1.ext", header_entries,
                                std::move(data));
 
     header_entries.insert(std::make_pair(kCustomField2, "CUSTOM2"));
     std::unique_ptr<std::string> data2(new std::string(kTestPayload2));
     upload_job->AddDataSegment("Name2", "", header_entries, std::move(data2));
     return upload_job;
   }
 
   content::TestBrowserThreadBundle test_browser_thread_bundle_;
   base::RunLoop run_loop_;
   net::EmbeddedTestServer test_server_;
   scoped_refptr<net::TestURLRequestContextGetter> request_context_getter_;
   MockOAuth2TokenService oauth2_service_;
 
   std::unique_ptr<UploadJob::ErrorCode> expected_error_;
+
+  // TaskRunner used to run individual tests.
+  scoped_refptr<ImmediateThreadTaskRunner> task_runner_;
 };
 
 class UploadFlowTest : public UploadJobTestBase {
  public:
   UploadFlowTest() {}
 
   // UploadJobTestBase:
   void SetUp() override {
     UploadJobTestBase::SetUp();
     test_server_.RegisterRequestHandler(
         base::Bind(&UploadFlowTest::HandlePostRequest, base::Unretained(this)));
+    upload_attempt_count_ = 0;
   }
 
   std::unique_ptr<net::test_server::HttpResponse> HandlePostRequest(
       const net::test_server::HttpRequest& request) {
+    upload_attempt_count_++;
     EXPECT_TRUE(request.headers.find("Authorization") != request.headers.end());
     const std::string authorization_header =
         request.headers.at("Authorization");
     std::unique_ptr<net::test_server::BasicHttpResponse> response(
         new net::test_server::BasicHttpResponse);
     const size_t pos = authorization_header.find(" ");
     if (pos == std::string::npos) {
       response->set_code(net::HTTP_UNAUTHORIZED);
       return std::move(response);
     }
 
     const std::string token = authorization_header.substr(pos + 1);
     response->set_code(oauth2_service_.IsTokenValid(token)
                            ? net::HTTP_OK
                            : net::HTTP_UNAUTHORIZED);
     return std::move(response);
   }
+
+ protected:
+  int upload_attempt_count_;
 };
 
 TEST_F(UploadFlowTest, SuccessfulUpload) {
   oauth2_service_.SetTokenValid(kTokenValid);
   oauth2_service_.AddTokenToQueue(kTokenValid);
   std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
       base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
   upload_job->Start();
   run_loop_.Run();
+  ASSERT_EQ(1, upload_attempt_count_);
 }
 
 TEST_F(UploadFlowTest, TokenExpired) {
   oauth2_service_.SetTokenValid(kTokenValid);
   oauth2_service_.AddTokenToQueue(kTokenExpired);
   oauth2_service_.AddTokenToQueue(kTokenValid);
   std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
       base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
   upload_job->Start();
   run_loop_.Run();
+  ASSERT_EQ(2, upload_attempt_count_);
 }
 
 TEST_F(UploadFlowTest, TokenInvalid) {
   oauth2_service_.AddTokenToQueue(kTokenInvalid);
   oauth2_service_.AddTokenToQueue(kTokenInvalid);
+  oauth2_service_.AddTokenToQueue(kTokenInvalid);
+  oauth2_service_.AddTokenToQueue(kTokenInvalid);
   SetExpectedError(std::unique_ptr<UploadJob::ErrorCode>(
       new UploadJob::ErrorCode(UploadJob::AUTHENTICATION_ERROR)));
 
   std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
       base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
   upload_job->Start();
   run_loop_.Run();
+  ASSERT_EQ(4, upload_attempt_count_);
+}
+
+TEST_F(UploadFlowTest, TokenMultipleTries) {
+  oauth2_service_.SetTokenValid(kTokenValid);
+  oauth2_service_.AddTokenToQueue(kTokenInvalid);
+  oauth2_service_.AddTokenToQueue(kTokenInvalid);
+  oauth2_service_.AddTokenToQueue(kTokenValid);
+
+  std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
+      base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
+  upload_job->Start();
+  run_loop_.Run();
+  ASSERT_EQ(3, upload_attempt_count_);
 }
 
 TEST_F(UploadFlowTest, TokenFetchFailure) {
   SetExpectedError(std::unique_ptr<UploadJob::ErrorCode>(
       new UploadJob::ErrorCode(UploadJob::AUTHENTICATION_ERROR)));
 
   std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
       base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
   upload_job->Start();
   run_loop_.Run();
+  // Without a token we don't try to upload
+  ASSERT_EQ(0, upload_attempt_count_);
+}
+
+class UploadServerErrorTest : public UploadJobTestBase {

[Andrew T Wilson (Slow), 2016/05/10 22:00:42]

Could you use UploadFlowTest() but have some member variable that makes
HandlePostRequest return INTERNAL_SERVER_ERROR, rather than making another test
class?

This is fine as-is, although it does have some code duplication between the two
classes it might be nice to avoid.

[Marton Hunyady, 2016/05/11 12:55:31]

Done, removed it.

+ public:
+  UploadServerErrorTest() {}
+
+  // UploadJobTestBase:
+  void SetUp() override {
+    UploadJobTestBase::SetUp();
+    test_server_.RegisterRequestHandler(base::Bind(
+        &UploadServerErrorTest::HandlePostRequest, base::Unretained(this)));
+    upload_attempt_count_ = 0;
+  }
+
+  std::unique_ptr<net::test_server::HttpResponse> HandlePostRequest(
+      const net::test_server::HttpRequest& request) {
+    upload_attempt_count_++;
+    std::unique_ptr<net::test_server::BasicHttpResponse> response(
+        new net::test_server::BasicHttpResponse);
+    response->set_code(net::HTTP_INTERNAL_SERVER_ERROR);
+    return std::move(response);
+  }
+
+ protected:
+  int upload_attempt_count_;
+};
+
+TEST_F(UploadServerErrorTest, InternalServerError) {
+  oauth2_service_.SetTokenValid(kTokenValid);
+  oauth2_service_.AddTokenToQueue(kTokenValid);
+
+  SetExpectedError(std::unique_ptr<UploadJob::ErrorCode>(
+      new UploadJob::ErrorCode(UploadJob::SERVER_ERROR)));
+
+  std::unique_ptr<UploadJob> upload_job = PrepareUploadJob(
+      base::WrapUnique(new UploadJobImpl::RandomMimeBoundaryGenerator));
+  upload_job->Start();
+  run_loop_.Run();
+
+  // kMaxAttempts
+  ASSERT_EQ(4, upload_attempt_count_);
 }
 
 class UploadRequestTest : public UploadJobTestBase {
  public:
   UploadRequestTest() {}
 
   // UploadJobTestBase:
   void SetUp() override {
     UploadJobTestBase::SetUp();
     test_server_.RegisterRequestHandler(base::Bind(
@@ skipping 35 matching lines @@
       "customfield2: CUSTOM2\r\n"
       "\r\n"
       "**||--||PAYLOAD2||--||**\r\n--"
       "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA--\r\n");
 
   upload_job->Start();
   run_loop_.Run();
 }
 
 }  // namespace policy