From ccf3adecc69e42c38fa2df3a8b76359e2d7784bd Mon Sep 17 00:00:00 2001
From: Isaac Murchie <isaac@saucelabs.com>
Date: Tue, 30 Jan 2018 14:16:27 -0500
Subject: [PATCH] Get chrome session back when resetting app

---
 lib/commands/general.js   |  2 +-
 lib/driver.js             | 15 ++++++++-------
 test/unit/driver-specs.js | 25 ++++++++++++++++++++++---
 3 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/lib/commands/general.js b/lib/commands/general.js
index ad0507e7..30c8dbe5 100644
--- a/lib/commands/general.js
+++ b/lib/commands/general.js
@@ -209,7 +209,7 @@ commands.reset = async function () {
   // reset context since we don't know what kind on context we will end up after app launch.
   this.curContext = NATIVE_WIN;
 
-  return await this.startAUT();
+  return await this.isChromeSession ? this.startChromeSession() : this.startAUT();
 };
 
 commands.startAUT = async function () {
diff --git a/lib/driver.js b/lib/driver.js
index fa073eaf..3d115fcc 100644
--- a/lib/driver.js
+++ b/lib/driver.js
@@ -2,7 +2,7 @@ import { BaseDriver, DeviceSettings } from 'appium-base-driver';
 import Chromedriver from 'appium-chromedriver';
 import desiredConstraints from './desired-caps';
 import commands from './commands/index';
-import { setupNewChromedriver } from './commands/context';
+import * as contextCommands from './commands/context';
 import helpers from './android-helpers';
 import { CHROMIUM_WIN } from './webview-helpers';
 import log from './logger';
@@ -261,10 +261,6 @@ class AndroidDriver extends BaseDriver {
     if (this.isChromeSession) {
       // start a chromedriver session and proxy to it
       await this.startChromeSession();
-      if (this.shouldDismissChromeWelcome()) {
-        // dismiss Chrome welcome dialog
-        await this.dismissChromeWelcome();
-      }
     } else {
       if (this.opts.autoLaunch) {
         // start app
@@ -365,8 +361,8 @@ class AndroidDriver extends BaseDriver {
     if (!_.includes(knownPackages, this.opts.appPackage)) {
       opts.chromeAndroidActivity = this.opts.appActivity;
     }
-    this.chromedriver = await setupNewChromedriver(opts, this.adb.curDeviceId,
-                                                   this.adb);
+    this.chromedriver = await contextCommands.setupNewChromedriver(opts, this.adb.curDeviceId,
+                                                                   this.adb);
     this.chromedriver.on(Chromedriver.EVENT_CHANGED, (msg) => {
       if (msg.state === Chromedriver.STATE_STOPPED) {
         this.onChromedriverStop(CHROMIUM_WIN);
@@ -380,6 +376,11 @@ class AndroidDriver extends BaseDriver {
     this.sessionChromedrivers[CHROMIUM_WIN] = this.chromedriver;
     this.proxyReqRes = this.chromedriver.proxyReq.bind(this.chromedriver);
     this.jwpProxyActive = true;
+
+    if (this.shouldDismissChromeWelcome()) {
+      // dismiss Chrome welcome dialog
+      await this.dismissChromeWelcome();
+    }
   }
 
   async checkAppPresent () {
diff --git a/test/unit/driver-specs.js b/test/unit/driver-specs.js
index 9a3897c6..e71d3340 100644
--- a/test/unit/driver-specs.js
+++ b/test/unit/driver-specs.js
@@ -3,13 +3,14 @@ import chaiAsPromised from 'chai-as-promised';
 import log from '../../lib/logger';
 import sinon from 'sinon';
 import helpers from '../../lib/android-helpers';
+import * as contextCommands from '../../lib/commands/context';
 import { withMocks } from 'appium-test-support';
 import AndroidDriver from '../..';
 import ADB from 'appium-adb';
 import { errors } from 'appium-base-driver';
 import { fs } from 'appium-support';
 import { SharedPrefsBuilder } from 'shared-preferences-builder';
-
+import _ from 'lodash';
 
 let driver;
 let sandbox = sinon.sandbox.create();
@@ -375,12 +376,30 @@ describe('driver', function () {
       await driver.startAndroidSession();
       driver.dismissChromeWelcome.calledOnce.should.be.false;
     });
+  });
+  describe('startChromeSession', function () {
+    beforeEach(async function () {
+      driver = new AndroidDriver();
+      driver.adb = new ADB();
+      driver.bootstrap = new helpers.bootstrap(driver.adb);
+      driver.settings = { update () { } };
+      driver.caps = {};
+
+      sandbox.stub(contextCommands, 'setupNewChromedriver').returns({
+        on: _.noop,
+        proxyReq: _.noop,
+      });
+      sandbox.stub(driver, 'dismissChromeWelcome');
+    });
+    afterEach(function () {
+      sandbox.restore();
+    });
     it('should call dismissChromeWelcome', async function () {
       driver.opts.browserName = 'Chrome';
       driver.opts.chromeOptions = {
-        "args" : ["--no-first-run"]
+        "args": ["--no-first-run"]
       };
-      await driver.startAndroidSession();
+      await driver.startChromeSession();
       driver.dismissChromeWelcome.calledOnce.should.be.true;
     });
   });