From 5a67708df02f153cab01c23e16f9f82d6d9c192c Mon Sep 17 00:00:00 2001
From: Isaac Murchie <isaac@saucelabs.com>
Date: Wed, 22 Jun 2016 08:54:02 -0700
Subject: [PATCH] Add ability to use automationName XCUITest

---
 lib/appium.js            | 12 +++++++++---
 lib/main.js              |  2 +-
 package.json             |  1 +
 test/driver-e2e-specs.js |  2 +-
 test/driver-specs.js     | 13 ++++++++++++-
 5 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/lib/appium.js b/lib/appium.js
index f39e7881e97..9eb21bc6e7f 100644
--- a/lib/appium.js
+++ b/lib/appium.js
@@ -7,6 +7,7 @@ import { FakeDriver } from 'appium-fake-driver';
 import { AndroidDriver } from 'appium-android-driver';
 import { IosDriver } from 'appium-ios-driver';
 import { SelendroidDriver } from 'appium-selendroid-driver';
+import { WebDriverAgentDriver } from 'appium-xcuitest-driver';
 import B from 'bluebird';
 import util from 'util';
 
@@ -40,9 +41,14 @@ class AppiumDriver extends BaseDriver {
     }
 
     // we don't necessarily have an `automationName` capability,
-    // but if we do and it is 'Selendroid', act on it
-    if ((caps.automationName || '').toLowerCase() === 'selendroid') {
-      return SelendroidDriver;
+    if (caps.automationName) {
+      if (caps.automationName.toLowerCase() === 'selendroid') {
+        // but if we do and it is 'Selendroid', act on it
+        return SelendroidDriver;
+      } else if (caps.automationName.toLowerCase() === 'xcuitest') {
+        // but if we do and it is 'XCUITest', act on it
+        return WebDriverAgentDriver;
+      }
     }
 
     if (caps.platformName.toLowerCase() === "fake") {
diff --git a/lib/main.js b/lib/main.js
index 4cf9dfee93a..fe8f32b9be4 100755
--- a/lib/main.js
+++ b/lib/main.js
@@ -110,7 +110,7 @@ async function main (args = null) {
       await registerNode(args.nodeconfig, args.address, args.port);
     }
   } catch (err) {
-    server.close();
+    await server.close();
     throw err;
   }
   logServerPort(args.address, args.port);
diff --git a/package.json b/package.json
index b647b487907..36ea0827163 100644
--- a/package.json
+++ b/package.json
@@ -42,6 +42,7 @@
     "appium-logger": "^2.1.0",
     "appium-selendroid-driver": "^1.3.4",
     "appium-support": "^2.3.0",
+    "appium-xcuitest-driver": "^1.0.4",
     "argparse": "^1.0.7",
     "asyncbox": "^2.3.1",
     "authorize-ios": "^1.0.4",
diff --git a/test/driver-e2e-specs.js b/test/driver-e2e-specs.js
index 0b94fc6b11c..fca58f77dc2 100644
--- a/test/driver-e2e-specs.js
+++ b/test/driver-e2e-specs.js
@@ -27,7 +27,7 @@ describe('FakeDriver - via HTTP', () => {
   });
   after(async () => {
     if (server) {
-      await B.promisify(server.close.bind(server))();
+      await server.close();
     }
   });
 
diff --git a/test/driver-specs.js b/test/driver-specs.js
index 3b9f5a18347..ae58d7eb289 100644
--- a/test/driver-specs.js
+++ b/test/driver-specs.js
@@ -7,6 +7,8 @@ import _ from 'lodash';
 import sinon from 'sinon';
 import chai from 'chai';
 import chaiAsPromised from 'chai-as-promised';
+import { WebDriverAgentDriver } from 'appium-xcuitest-driver';
+
 
 chai.use(chaiAsPromised);
 
@@ -170,10 +172,19 @@ describe('AppiumDriver', () => {
     describe('sessionExists', () => {
     });
     describe('getDriverForCaps', () => {
-      it('should not blow up if user doesnt provide platformName', () => {
+      it('should not blow up if user does not provide platformName', () => {
         let appium = new AppiumDriver({});
         (() => { appium.getDriverForCaps({}); }).should.throw(/platformName/);
       });
+      it('should get WebDriverAgentDriver driver for automationName of XCUITest', () => {
+        let appium = new AppiumDriver({});
+        let driver = appium.getDriverForCaps({
+          platformName: 'iOS',
+          automationName: 'XCUITest'
+        });
+        driver.should.be.an.instanceof(Function);
+        driver.should.equal(WebDriverAgentDriver);
+      });
     });
   });
 });