diff --git a/lib/main.js b/lib/main.js
index fe6551ec..4bbacb9d 100644
--- a/lib/main.js
+++ b/lib/main.js
@@ -420,12 +420,12 @@ Lambda.prototype._setRunTimeEnvironmentVars = function (program) {
   }
 }
 
-Lambda.prototype._uploadExisting = function (lambda, params, cb) {
-  var request = lambda.updateFunctionCode({
+Lambda.prototype._uploadExisting = (lambda, params, cb) => {
+  const request = lambda.updateFunctionCode({
     'FunctionName': params.FunctionName,
     'ZipFile': params.Code.ZipFile,
     'Publish': params.Publish
-  }, function (err, data) {
+  }, (err, data) => {
     if (err) {
       return cb(err, data)
     }
@@ -442,12 +442,12 @@ Lambda.prototype._uploadExisting = function (lambda, params, cb) {
       'Environment': params.Environment,
       'DeadLetterConfig': params.DeadLetterConfig,
       'TracingConfig': params.TracingConfig
-    }, function (err, data) {
+    }, (err, data) => {
       return cb(err, data)
     })
   })
 
-  request.on('retry', function (response) {
+  request.on('retry', (response) => {
     console.log(response.error.message)
     console.log('=> Retrying')
   })
@@ -455,12 +455,9 @@ Lambda.prototype._uploadExisting = function (lambda, params, cb) {
   return request
 }
 
-Lambda.prototype._uploadNew = function (lambda, params, cb) {
-  var request = lambda.createFunction(params, function (err, data) {
-    return cb(err, data)
-  })
-
-  request.on('retry', function (response) {
+Lambda.prototype._uploadNew = (lambda, params, cb) => {
+  const request = lambda.createFunction(params, (err, data) => cb(err, data))
+  request.on('retry', (response) => {
     console.log(response.error.message)
     console.log('=> Retrying')
   })
diff --git a/test/main.js b/test/main.js
index 6b9bb454..a34a235f 100644
--- a/test/main.js
+++ b/test/main.js
@@ -95,6 +95,13 @@ const _mockSetting = () => {
       callback(null, lambdaMockSettings[method])
     })
   })
+
+  return require('aws-sdk')
+}
+
+const _awsRestore = () => {
+  awsMock.restore('CloudWatchEvents')
+  awsMock.restore('Lambda')
 }
 
 /* global before, after, beforeEach, afterEach, describe, it */
@@ -105,6 +112,14 @@ describe('lib/main', function () {
     this.timeout(60000)
   }
 
+  let aws = null // mock
+  let awsLambda = null // mock
+  before(() => {
+    aws = _mockSetting()
+    awsLambda = new aws.Lambda({ apiVersion: '2015-03-31' })
+  })
+  after(() => _awsRestore())
+
   beforeEach(function () {
     program = Hoek.clone(originalProgram)
   })
@@ -834,25 +849,14 @@ describe('lib/main', function () {
       }]
     }
 
-    let awsLambda = null
-
     before(() => {
       fs.writeFileSync(
         'event_sources.json',
         JSON.stringify(eventSourcesJsonValue)
       )
-      _mockSetting()
-
-      awsLambda = new (require('aws-sdk')).Lambda({
-        apiVersion: '2015-03-31'
-      })
     })
 
-    after(() => {
-      fs.unlinkSync('event_sources.json')
-      awsMock.restore('CloudWatchEvents')
-      awsMock.restore('Lambda')
-    })
+    after(() => fs.unlinkSync('event_sources.json'))
 
     it('program.eventSourceFile is empty value', (done) => {
       program.eventSourceFile = ''
@@ -935,16 +939,10 @@ describe('lib/main', function () {
         'event_sources.json',
         JSON.stringify(eventSourcesJsonValue)
       )
-
-      _mockSetting()
-      schedule = new ScheduleEvents(require('aws-sdk'))
+      schedule = new ScheduleEvents(aws)
     })
 
-    after(() => {
-      fs.unlinkSync('event_sources.json')
-      awsMock.restore('CloudWatchEvents')
-      awsMock.restore('Lambda')
-    })
+    after(() => fs.unlinkSync('event_sources.json'))
 
     it('program.eventSourceFile is empty value', (done) => {
       program.eventSourceFile = ''
@@ -985,6 +983,28 @@ describe('lib/main', function () {
     })
   })
 
+  describe('_uploadNew', () => {
+    it('simple test with mock', (done) => {
+      const params = lambda._params(program, null)
+      lambda._uploadNew(awsLambda, params, (err, results) => {
+        assert.isNull(err)
+        assert.deepEqual(results, lambdaMockSettings.createFunction)
+        done()
+      })
+    })
+  })
+
+  describe('_uploadExisting', () => {
+    it('simple test with mock', (done) => {
+      const params = lambda._params(program, null)
+      lambda._uploadExisting(awsLambda, params, (err, results) => {
+        assert.isNull(err)
+        assert.deepEqual(results, lambdaMockSettings.updateFunctionConfiguration)
+        done()
+      })
+    })
+  })
+
   describe('check env vars before create sample files', function () {
     const filesCreatedBySetup = [
       '.env',