From 65bec1bddbee9e5b0da6d51c72f3e3eaf9be2c19 Mon Sep 17 00:00:00 2001
From: Patrick Hulce <patrick.hulce@gmail.com>
Date: Tue, 21 Mar 2017 19:50:49 -0700
Subject: [PATCH] feat(cli): add support for custom trace categories (#1866)

---
 lighthouse-cli/bin.ts                      |  2 ++
 lighthouse-core/gather/driver.js           | 12 ++++++--
 lighthouse-core/gather/gather-runner.js    |  2 +-
 lighthouse-core/test/gather/driver-test.js | 32 ++++++++++++++++++++++
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/lighthouse-cli/bin.ts b/lighthouse-cli/bin.ts
index 6e8b4598b423..ae9a89de29a3 100755
--- a/lighthouse-cli/bin.ts
+++ b/lighthouse-cli/bin.ts
@@ -63,6 +63,7 @@ const cliFlags = yargs
     'save-artifacts',
     'list-all-audits',
     'list-trace-categories',
+    'additional-trace-categories',
     'config-path',
     'chrome-flags',
     'perf',
@@ -78,6 +79,7 @@ const cliFlags = yargs
     'save-artifacts': 'Save all gathered artifacts to disk',
     'list-all-audits': 'Prints a list of all available audits and exits',
     'list-trace-categories': 'Prints a list of all required trace categories and exits',
+    'additional-trace-categories': 'Additional categories to capture with the trace (comma-delimited).',
     'config-path': 'The path to the config JSON.',
     'chrome-flags': 'Custom flags to pass to Chrome.',
     'perf': 'Use a performance-test-only configuration',
diff --git a/lighthouse-core/gather/driver.js b/lighthouse-core/gather/driver.js
index 2843f360f6c7..9bec04966216 100644
--- a/lighthouse-core/gather/driver.js
+++ b/lighthouse-core/gather/driver.js
@@ -27,6 +27,8 @@ const DevtoolsLog = require('./devtools-log');
 
 const PAUSE_AFTER_LOAD = 500;
 
+const _uniq = arr => Array.from(new Set(arr));
+
 class Driver {
   static get MAX_WAIT_FOR_FULLY_LOADED() {
     return 25 * 1000;
@@ -570,9 +572,15 @@ class Driver {
       });
   }
 
-  beginTrace() {
+  /**
+   * @param {{additionalTraceCategories: string=}=} flags
+   */
+  beginTrace(flags) {
+    const additionalCategories = (flags && flags.additionalTraceCategories &&
+        flags.additionalTraceCategories.split(',')) || [];
+    const traceCategories = this._traceCategories.concat(additionalCategories);
     const tracingOpts = {
-      categories: this._traceCategories.join(','),
+      categories: _uniq(traceCategories).join(','),
       transferMode: 'ReturnAsStream',
       options: 'sampling-frequency=10000'  // 1000 is default and too slow.
     };
diff --git a/lighthouse-core/gather/gather-runner.js b/lighthouse-core/gather/gather-runner.js
index 4ab2771d445a..494844156f9b 100644
--- a/lighthouse-core/gather/gather-runner.js
+++ b/lighthouse-core/gather/gather-runner.js
@@ -80,7 +80,7 @@ class GatherRunner {
   static loadPage(driver, options) {
     return Promise.resolve()
       // Begin tracing only if requested by config.
-      .then(_ => options.config.recordTrace && driver.beginTrace())
+      .then(_ => options.config.recordTrace && driver.beginTrace(options.flags))
       // Network is always recorded for internal use, even if not saved as artifact.
       .then(_ => driver.beginNetworkCollect(options))
       // Navigate.
diff --git a/lighthouse-core/test/gather/driver-test.js b/lighthouse-core/test/gather/driver-test.js
index 39dc44f51de3..aa1bc85fe588 100644
--- a/lighthouse-core/test/gather/driver-test.js
+++ b/lighthouse-core/test/gather/driver-test.js
@@ -17,6 +17,8 @@
 
 'use strict';
 
+let sendCommandParams = [];
+
 const Driver = require('../../gather/driver.js');
 const Connection = require('../../gather/connections/connection.js');
 const Element = require('../../lib/element.js');
@@ -53,6 +55,7 @@ function createActiveWorker(id, url, controlledClients) {
 }
 
 connection.sendCommand = function(command, params) {
+  sendCommandParams.push({command, params});
   switch (command) {
     case 'DOM.getDocument':
       return Promise.resolve({root: {nodeId: 249}});
@@ -73,6 +76,8 @@ connection.sendCommand = function(command, params) {
           }
         }]
       });
+    case 'Page.enable':
+    case 'Tracing.start':
     case 'ServiceWorker.enable':
     case 'ServiceWorker.disable':
       return Promise.resolve();
@@ -98,6 +103,10 @@ const mockRedirects = [req1, req2, req3];
 /* eslint-env mocha */
 
 describe('Browser Driver', () => {
+  beforeEach(() => {
+    sendCommandParams = [];
+  });
+
   it('returns null when DOM.querySelector finds no node', () => {
     return driverStub.querySelector('invalid').then(value => {
       assert.equal(value, null);
@@ -150,9 +159,32 @@ describe('Browser Driver', () => {
     assert.notEqual(opts.url, req1.url, 'opts.url changed after the redirects');
     assert.equal(opts.url, req3.url, 'opts.url matches the last redirect');
   });
+
+  it('will request default traceCategories', () => {
+    return driverStub.beginTrace().then(() => {
+      const traceCmd = sendCommandParams.find(obj => obj.command === 'Tracing.start');
+      const categories = traceCmd.params.categories;
+      assert.ok(categories.includes('devtools.timeline'), 'contains devtools.timeline');
+    });
+  });
+
+  it('will use requested additionalTraceCategories', () => {
+    return driverStub.beginTrace({additionalTraceCategories: 'v8,v8.execute,toplevel'}).then(() => {
+      const traceCmd = sendCommandParams.find(obj => obj.command === 'Tracing.start');
+      const categories = traceCmd.params.categories;
+      assert.ok(categories.includes('blink'), 'contains default categories');
+      assert.ok(categories.includes('v8.execute'), 'contains added categories');
+      assert.ok(categories.indexOf('toplevel') === categories.lastIndexOf('toplevel'),
+          'de-dupes categories');
+    });
+  });
 });
 
 describe('Multiple tab check', () => {
+  beforeEach(() => {
+    sendCommandParams = [];
+  });
+
   it('will pass if there are no current service workers', () => {
     const pageUrl = 'https://example.com/';
     driverStub.once = createOnceStub({