Adds a --deterministic-fetch flag to headless_shell

headless/app/headless_shell.cc

 // Copyright 2016 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 <iostream>
 #include <memory>
 #include <string>
 
 #include "base/base64.h"
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/command_line.h"
 #include "base/files/file_path.h"
 #include "base/json/json_writer.h"
 #include "base/location.h"
 #include "base/memory/ref_counted.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/strings/string_number_conversions.h"
 #include "content/public/common/content_switches.h"
 #include "headless/app/headless_shell_switches.h"
 #include "headless/public/domains/emulation.h"
 #include "headless/public/domains/page.h"
 #include "headless/public/domains/runtime.h"
 #include "headless/public/headless_browser.h"
 #include "headless/public/headless_devtools_client.h"
 #include "headless/public/headless_devtools_target.h"
 #include "headless/public/headless_web_contents.h"
+#include "headless/public/util/deterministic_dispatcher.h"
+#include "headless/public/util/deterministic_http_protocol_handler.h"
 #include "net/base/file_stream.h"
 #include "net/base/io_buffer.h"
 #include "net/base/ip_address.h"
 #include "net/base/net_errors.h"
 #include "ui/gfx/geometry/size.h"
 
 using headless::HeadlessBrowser;
 using headless::HeadlessBrowserContext;
 using headless::HeadlessDevToolsClient;
 using headless::HeadlessWebContents;
@@ skipping 28 matching lines @@
       : browser_(nullptr),
         devtools_client_(HeadlessDevToolsClient::Create()),
         web_contents_(nullptr),
         processed_page_ready_(false),
         browser_context_(nullptr) {}
   ~HeadlessShell() override {}
 
   void OnStart(HeadlessBrowser* browser) {
     browser_ = browser;
 
-    browser_context_ = browser_->CreateBrowserContextBuilder().Build();
+    if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+            headless::switches::kDeterministicFetch)) {
+      deterministic_dispatcher_.reset(
+          new headless::DeterministicDispatcher(browser_->BrowserIOThread()));
+
+      headless::ProtocolHandlerMap protocol_handlers;
+      protocol_handlers[url::kHttpScheme] =
+          base::MakeUnique<headless::DeterministicHttpProtocolHandler>(
+              deterministic_dispatcher_.get(), browser->BrowserIOThread());
+      protocol_handlers[url::kHttpsScheme] =
+          base::MakeUnique<headless::DeterministicHttpProtocolHandler>(
+              deterministic_dispatcher_.get(), browser->BrowserIOThread());
+
+      browser_context_ = browser_->CreateBrowserContextBuilder()
+                             .SetProtocolHandlers(std::move(protocol_handlers))
+                             .Build();
+    } else {
+      browser_context_ = browser_->CreateBrowserContextBuilder().Build();
+    }
 
     HeadlessWebContents::Builder builder(
         browser_context_->CreateWebContentsBuilder());
     base::CommandLine::StringVector args =
         base::CommandLine::ForCurrentProcess()->GetArgs();
 
     // TODO(alexclarke): Should we navigate to about:blank first if using
     // virtual time?
     if (!args.empty() && !args[0].empty())
       builder.SetInitialURL(GURL(args[0]));
@@ skipping 241 matching lines @@
   }
 
  private:
   GURL url_;
   HeadlessBrowser* browser_;  // Not owned.
   std::unique_ptr<HeadlessDevToolsClient> devtools_client_;
   HeadlessWebContents* web_contents_;
   bool processed_page_ready_;
   std::unique_ptr<net::FileStream> screenshot_file_stream_;
   HeadlessBrowserContext* browser_context_;
+  std::unique_ptr<headless::DeterministicDispatcher> deterministic_dispatcher_;
 
   DISALLOW_COPY_AND_ASSIGN(HeadlessShell);
 };
 
 int main(int argc, const char** argv) {
   headless::RunChildProcessIfNeeded(argc, argv);
   HeadlessShell shell;
   HeadlessBrowser::Options::Builder builder(argc, argv);
 
   // Enable devtools if requested.
@@ skipping 60 matching lines @@
       LOG(ERROR) << "Malformed window size";
       return EXIT_FAILURE;
     }
     builder.SetWindowSize(parsed_window_size);
   }
 
   return HeadlessBrowserMain(
       builder.Build(),
       base::Bind(&HeadlessShell::OnStart, base::Unretained(&shell)));
 }