Expand OAuth2ApiCallFlow class to allow multiple lines of headers & c…

…ustomizable expected http response code

BUG=1154032

Change-Id: Ia92533a9e8d0fa08d1156940d232314633057f08
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2577004
Commit-Queue: Alice Gong <alicego@google.com>
Reviewed-by: Alex Ilin <alexilin@chromium.org>
Cr-Commit-Position: refs/heads/master@{#834344}

google_apis/gaia/oauth2_api_call_flow.cc

@@ -13,6 +13,7 @@
 #include "google_apis/gaia/gaia_urls.h"
 #include "net/base/escape.h"
 #include "net/base/load_flags.h"
+#include "net/http/http_request_headers.h"
 #include "net/http/http_status_code.h"
 #include "services/network/public/cpp/resource_request.h"
 #include "services/network/public/cpp/shared_url_loader_factory.h"
@@ -30,7 +31,7 @@ static std::string MakeAuthorizationValue(const std::string& auth_token) {
 OAuth2ApiCallFlow::OAuth2ApiCallFlow() : state_(INITIAL) {
 }
 
-OAuth2ApiCallFlow::~OAuth2ApiCallFlow() {}
+OAuth2ApiCallFlow::~OAuth2ApiCallFlow() = default;
 
 void OAuth2ApiCallFlow::Start(
     scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
@@ -45,15 +46,18 @@ void OAuth2ApiCallFlow::Start(
                      base::Unretained(this)));
 }
 
+net::HttpRequestHeaders OAuth2ApiCallFlow::CreateApiCallHeaders() {
+  return net::HttpRequestHeaders();
+}
+
 void OAuth2ApiCallFlow::EndApiCall(std::unique_ptr<std::string> body) {
   CHECK_EQ(API_CALL_STARTED, state_);
   std::unique_ptr<network::SimpleURLLoader> source = std::move(url_loader_);
 
   int status_code = 0;
   if (source->ResponseInfo() && source->ResponseInfo()->headers)
     status_code = source->ResponseInfo()->headers->response_code();
-  if (source->NetError() != net::OK ||
-      (status_code != net::HTTP_OK && status_code != net::HTTP_NO_CONTENT)) {
+  if (source->NetError() != net::OK || !IsExpectedSuccessCode(status_code)) {
     state_ = ERROR_STATE;
     ProcessApiCallFailure(source->NetError(), source->ResponseInfo(),
                           std::move(body));
@@ -71,6 +75,10 @@ std::string OAuth2ApiCallFlow::GetRequestTypeForBody(const std::string& body) {
   return body.empty() ? "GET" : "POST";
 }
 
+bool OAuth2ApiCallFlow::IsExpectedSuccessCode(int code) const {
+  return code == net::HTTP_OK || code == net::HTTP_NO_CONTENT;
+}
+
 void OAuth2ApiCallFlow::OnURLLoadComplete(std::unique_ptr<std::string> body) {
   CHECK_EQ(API_CALL_STARTED, state_);
   EndApiCall(std::move(body));
@@ -91,8 +99,10 @@ std::unique_ptr<network::SimpleURLLoader> OAuth2ApiCallFlow::CreateURLLoader(
   request->url = CreateApiCallUrl();
   request->method = request_type;
   request->credentials_mode = network::mojom::CredentialsMode::kOmit;
+  request->headers = CreateApiCallHeaders();
   request->headers.SetHeader("Authorization",
                              MakeAuthorizationValue(access_token));
+
   std::unique_ptr<network::SimpleURLLoader> result =
       network::SimpleURLLoader::Create(std::move(request), traffic_annotation);
 

google_apis/gaia/oauth2_api_call_flow.h

@@ -19,6 +19,10 @@ class SimpleURLLoader;
 class SharedURLLoaderFactory;
 }
 
+namespace net {
+class HttpRequestHeaders;
+}
+
 // Base class for all classes that implement a flow to call OAuth2 enabled APIs,
 // given an access token to the service.  This class abstracts the basic steps
 // and exposes template methods for sub-classes to implement for API specific
@@ -39,13 +43,19 @@ class OAuth2ApiCallFlow {
 
   // Methods to help create the API request.
   virtual GURL CreateApiCallUrl() = 0;
+  virtual net::HttpRequestHeaders CreateApiCallHeaders();
   virtual std::string CreateApiCallBody() = 0;
   virtual std::string CreateApiCallBodyContentType();
 
   // Returns the request type (e.g. GET, POST) for the |body| that will be sent
   // with the request.
   virtual std::string GetRequestTypeForBody(const std::string& body);
 
+  // Called when the API call ends to check whether it succeeded, and decide
+  // which of the following 2 process functions to call. Should be overriden by
+  // subclasses if the expected success response code is not 200 or 204.
+  virtual bool IsExpectedSuccessCode(int code) const;
+
   // Sub-classes can expose an appropriate observer interface by implementing
   // these template methods.
   // Called when the API call finished successfully. |body| may be null.

google_apis/gaia/oauth2_api_call_flow_unittest.cc

@@ -54,6 +54,7 @@ class MockApiCallFlow : public OAuth2ApiCallFlow {
 
   MOCK_METHOD0(CreateApiCallUrl, GURL());
   MOCK_METHOD0(CreateApiCallBody, std::string());
+  MOCK_METHOD0(CreateApiCallHeaders, net::HttpRequestHeaders());
   MOCK_METHOD2(ProcessApiCallSuccess,
                void(const network::mojom::URLResponseHead* head,
                     std::unique_ptr<std::string> body));
@@ -92,11 +93,13 @@ class OAuth2ApiCallFlowTest : public testing::Test {
         network::URLLoaderCompletionStatus(error));
   }
 
+ protected:
   void SetupApiCall(bool succeeds, net::HttpStatusCode status) {
     std::string body(CreateBody());
     GURL url(CreateApiUrl());
     EXPECT_CALL(flow_, CreateApiCallBody()).WillOnce(Return(body));
     EXPECT_CALL(flow_, CreateApiCallUrl()).WillOnce(Return(url));
+    EXPECT_CALL(flow_, CreateApiCallHeaders());
 
     AddFetchResult(url, succeeds, status, std::string());
   }
@@ -150,3 +153,39 @@ TEST_F(OAuth2ApiCallFlowTest, ExpectedHTTPHeaders) {
   EXPECT_EQ("Bearer access_token", auth_header);
   EXPECT_EQ(body, network::GetUploadData(pending[0].request));
 }
+
+net::HttpRequestHeaders CreateHeaders() {
+  net::HttpRequestHeaders headers;
+  headers.SetHeader("Test-Header-Field", "test content");
+  return headers;
+}
+
+TEST_F(OAuth2ApiCallFlowTest, ExpectedMultipleHTTPHeaders) {
+  std::string body = CreateBody();
+  GURL url(CreateApiUrl());
+  SetupApiCall(true, net::HTTP_OK);
+
+  // Overwrite EXPECT_CALL default return so that we get multiple headers.
+  ON_CALL(flow_, CreateApiCallHeaders).WillByDefault(Return(CreateHeaders()));
+
+  // ... never mind the HTTP response part of the setup --- don't want
+  // TestURLLoaderFactory replying to it just yet as it would prevent examining
+  // the request headers.
+  test_url_loader_factory_.ClearResponses();
+
+  flow_.Start(shared_factory_, kAccessToken);
+  const std::vector<network::TestURLLoaderFactory::PendingRequest>& pending =
+      *test_url_loader_factory_.pending_requests();
+  ASSERT_EQ(1u, pending.size());
+  EXPECT_EQ(url, pending[0].request.url);
+
+  const auto& headers = pending[0].request.headers;
+  std::string auth_header;
+  EXPECT_TRUE(headers.GetHeader("Authorization", &auth_header));
+  EXPECT_EQ("Bearer access_token", auth_header);
+  std::string test_header_content;
+  EXPECT_TRUE(headers.GetHeader("Test-Header-Field", &test_header_content));
+  EXPECT_EQ("test content", test_header_content);
+
+  EXPECT_EQ(body, network::GetUploadData(pending[0].request));
+}