Add a --virtual-time-budget option 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 "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;
+namespace emulation = headless::emulation;
 namespace page = headless::page;
 namespace runtime = headless::runtime;
 
 namespace {
 // Address where to listen to incoming DevTools connections.
 const char kDevToolsHttpServerAddress[] = "127.0.0.1";
 // Default file name for screenshot. Can be overriden by "--screenshot" switch.
 const char kDefaultScreenshotFileName[] = "screenshot.png";
 
 bool ParseWindowSize(std::string window_size, gfx::Size* parsed_window_size) {
   int width, height = 0;
   if (sscanf(window_size.c_str(), "%dx%d", &width, &height) >= 2 &&
       width >= 0 && height >= 0) {
     parsed_window_size->set_width(width);
     parsed_window_size->set_height(height);
     return true;
   }
   return false;
 }
 }  // namespace
 
 // A sample application which demonstrates the use of the headless API.
-class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
+class HeadlessShell : public HeadlessWebContents::Observer,
+                      emulation::ExperimentalObserver,
+                      page::Observer {
  public:
   HeadlessShell()
       : 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();
 
     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]));
 
     web_contents_ = builder.Build();
     if (!web_contents_) {
       LOG(ERROR) << "Navigation failed";
       browser_->Shutdown();
       return;
     }
     web_contents_->AddObserver(this);
   }
 
   void Shutdown() {
     if (!web_contents_)
       return;
     if (!RemoteDebuggingEnabled()) {
+      devtools_client_->GetEmulation()->GetExperimental()->RemoveObserver(this);
       devtools_client_->GetPage()->RemoveObserver(this);
       web_contents_->GetDevToolsTarget()->DetachClient(devtools_client_.get());
     }
     web_contents_->RemoveObserver(this);
     web_contents_ = nullptr;
     browser_context_->Close();
     browser_->Shutdown();
   }
 
   // HeadlessWebContents::Observer implementation:
   void DevToolsTargetReady() override {
     if (RemoteDebuggingEnabled())
       return;
     web_contents_->GetDevToolsTarget()->AttachClient(devtools_client_.get());
     devtools_client_->GetPage()->AddObserver(this);
     devtools_client_->GetPage()->Enable();
     // Check if the document had already finished loading by the time we
     // attached.
-    PollReadyState();
+
+    devtools_client_->GetEmulation()->GetExperimental()->AddObserver(this);
+
+    if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+            headless::switches::kVirtualTime)) {
+      devtools_client_->GetEmulation()->GetExperimental()->SetVirtualTimePolicy(
+          emulation::SetVirtualTimePolicyParams::Builder()
+              .SetPolicy(emulation::VirtualTimePolicy::
+                             PAUSE_IF_NETWORK_FETCHES_PENDING)
+              .SetBudget(10000)  // 10 virtual seconds

[Sami, 2016/08/08 10:12:28]

Mind making this configurable so that you could say --virtual-time=NN?

[alex clarke (OOO till 29th), 2016/08/08 20:28:12]

Done.

+              .Build());
+    } else {
+      PollReadyState();
+    }
     // TODO(skyostil): Implement more features to demonstrate the devtools API.
   }
 
   void PollReadyState() {
     // We need to check the current location in addition to the ready state to
     // be sure the expected page is ready.
     devtools_client_->GetRuntime()->Evaluate(
         "document.readyState + ' ' + document.location.href",
         base::Bind(&HeadlessShell::OnReadyState, base::Unretained(this)));
   }
 
   void OnReadyState(std::unique_ptr<runtime::EvaluateResult> result) {
     std::string ready_state_and_url;
     if (result->GetResult()->GetValue()->GetAsString(&ready_state_and_url)) {
       std::stringstream stream(ready_state_and_url);
       std::string ready_state;
       std::string url;
       stream >> ready_state;
       stream >> url;
 
       if (ready_state == "complete" &&
           (url_.spec() == url || url != "about:blank")) {
         OnPageReady();
         return;
       }
     }
   }
 
+  // emulation::Observer implementation:
+  void OnVirtualTimeBudgetExpired(
+      const emulation::VirtualTimeBudgetExpiredParams& params) override {
+    OnPageReady();
+  }
+
   // page::Observer implementation:
   void OnLoadEventFired(const page::LoadEventFiredParams& params) override {
+    if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+            headless::switches::kVirtualTime)) {
+      return;
+    }
     OnPageReady();
   }
 
   void OnPageReady() {
     if (processed_page_ready_)
       return;
     processed_page_ready_ = true;
 
     if (base::CommandLine::ForCurrentProcess()->HasSwitch(
             headless::switches::kDumpDom)) {
@@ skipping 217 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)));
 }