diff --git a/.babelrc b/.babelrc
index 9d8d516..e5d85e0 100644
--- a/.babelrc
+++ b/.babelrc
@@ -1 +1,24 @@
-{ "presets": ["es2015"] }
+{
+    "plugins": [
+        "check-es2015-constants",
+        "transform-es2015-arrow-functions",
+        "transform-es2015-block-scoped-functions",
+        "transform-es2015-block-scoping",
+        "transform-es2015-classes",
+        "transform-es2015-computed-properties",
+        "transform-es2015-destructuring",
+        "transform-es2015-duplicate-keys",
+        "transform-es2015-for-of",
+        "transform-es2015-function-name",
+        "transform-es2015-literals",
+        "transform-es2015-modules-commonjs",
+        "transform-es2015-object-super",
+        "transform-es2015-parameters",
+        "transform-es2015-shorthand-properties",
+        "transform-es2015-spread",
+        "transform-es2015-sticky-regex",
+        "transform-es2015-template-literals",
+        "transform-es2015-typeof-symbol",
+        "transform-es2015-unicode-regex"
+    ]
+}
diff --git a/gruntfile.js b/gruntfile.js
index be9bf99..f55dfe0 100644
--- a/gruntfile.js
+++ b/gruntfile.js
@@ -11,7 +11,7 @@ module.exports = function (grunt) {
                 files: [{
                     expand: true,
                     cwd: './lib',
-                    src: ['*.js'],
+                    src: ['**/*.js'],
                     dest: 'build',
                     ext: '.js'
                 }]
@@ -55,7 +55,6 @@ module.exports = function (grunt) {
     })
 
     require('load-grunt-tasks')(grunt)
-    grunt.loadTasks('build')
     grunt.registerTask('default', ['build'])
     grunt.registerTask('build', 'Build grunt-webdriver', function () {
         grunt.task.run([
diff --git a/lib/documentScreenshot.js b/lib/documentScreenshot.js
index 93cd926..59312c6 100644
--- a/lib/documentScreenshot.js
+++ b/lib/documentScreenshot.js
@@ -31,7 +31,7 @@ import getPageInfo from './scripts/getPageInfo'
 
 const CAPS_TAKING_FULLSIZE_SHOTS = ['firefox']
 
-export async function documentScreenshot (fileName) {
+export default async function documentScreenshot (fileName) {
     let ErrorHandler = this.instance.ErrorHandler
 
     /*!
diff --git a/lib/viewportScreenshot.js b/lib/viewportScreenshot.js
index cda4dd4..08f8e52 100644
--- a/lib/viewportScreenshot.js
+++ b/lib/viewportScreenshot.js
@@ -5,7 +5,7 @@
  */
 
 import gm from 'gm'
-import ErrorHandler from 'webdriverio/lib/utils/ErrorHandler'
+import { ErrorHandler } from 'webdriverio'
 
 import getScrollingPosition from './scripts/getScrollingPosition'
 import getScreenDimension from './scripts/getScreenDimension'
diff --git a/lib/webdrivercss.js b/lib/webdrivercss.js
index 585eed4..03c46a3 100644
--- a/lib/webdrivercss.js
+++ b/lib/webdrivercss.js
@@ -13,7 +13,8 @@ const DEFAULT_PROPERTIES = {
     screenWidth: [],
     warning: [],
     resultObject: {},
-    updateBaseline: false
+    updateBaseline: false,
+    applitools: {}
 }
 
 /**
diff --git a/test/bootstrap.js b/test/bootstrap.js
index 335921a..c675c72 100644
--- a/test/bootstrap.js
+++ b/test/bootstrap.js
@@ -31,7 +31,11 @@ afterHook = function(done) {
      * close browser and clean up created directories
      */
     async.parallel([
-        function(done) { browser.end(done); },
+        function(done) {
+            browser.end()
+                .then(done.bind(null, null))
+                .catch(done);
+        },
         function(done) { fs.remove(failedComparisonsRootDefault,done); },
         function(done) { fs.remove(screenshotRootDefault,done); },
         function(done) { fs.remove(failedComparisonsRootCustom,done); },
diff --git a/test/spec/instantiation.js b/test/spec/instantiation.js
index 9302df2..be13c35 100644
--- a/test/spec/instantiation.js
+++ b/test/spec/instantiation.js
@@ -2,8 +2,8 @@
 
 describe('WebdriverCSS plugin as WebdriverIO enhancement', function() {
 
-    before(function(done) {
-        this.browser = WebdriverIO.remote(capabilities).init(done);
+    before(function() {
+        return this.browser = WebdriverIO.remote(capabilities).init();
     });
 
     it('should not exist as command in WebdriverIO instance without initialization', function() {
@@ -32,9 +32,9 @@ describe('WebdriverCSS plugin as WebdriverIO enhancement', function() {
     it('should contain some default values', function() {
         var plugin = WebdriverCSS.init(this.browser);
 
-        expect(plugin).to.have.property('screenshotRoot').to.equal('webdrivercss');
-        expect(plugin).to.have.property('failedComparisonsRoot').to.equal('webdrivercss/diff');
-        expect(plugin).to.have.property('misMatchTolerance').to.equal(0.05);
+        expect(plugin.options).to.have.property('screenshotRoot').to.equal('webdrivercss');
+        expect(plugin.options).to.have.property('failedComparisonsRoot').to.equal('webdrivercss/diff');
+        expect(plugin.options).to.have.property('misMatchTolerance').to.equal(0.05);
 
     });
 
@@ -45,9 +45,9 @@ describe('WebdriverCSS plugin as WebdriverIO enhancement', function() {
             misMatchTolerance: 50
         });
 
-        expect(plugin).to.have.property('screenshotRoot').to.equal('__screenshotRoot__');
-        expect(plugin).to.have.property('failedComparisonsRoot').to.equal('__failedComparisonsRoot__');
-        expect(plugin).to.have.property('misMatchTolerance').to.equal(50);
+        expect(plugin.options).to.have.property('screenshotRoot').to.equal('__screenshotRoot__');
+        expect(plugin.options).to.have.property('failedComparisonsRoot').to.equal('__failedComparisonsRoot__');
+        expect(plugin.options).to.have.property('misMatchTolerance').to.equal(50);
     });
 
     it('should have a created "screenshotRoot" folder after initialization', function(done) {