From 3a210f602b79f3982094da6f2e35fcb90a5a956d Mon Sep 17 00:00:00 2001
From: Anthony Laibe <anthony@laibe.cc>
Date: Mon, 17 Sep 2018 09:57:51 +0100
Subject: [PATCH 1/2] Better support for await

---
 lib/core/modules/coderunner/codeRunner.js | 32 +++++++++++++++--------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/lib/core/modules/coderunner/codeRunner.js b/lib/core/modules/coderunner/codeRunner.js
index 167e23f65..a00db4b10 100644
--- a/lib/core/modules/coderunner/codeRunner.js
+++ b/lib/core/modules/coderunner/codeRunner.js
@@ -1,4 +1,5 @@
 const RunCode = require('./runCode.js');
+const Utils = require('../../../utils/utils');
 
 class CodeRunner {
   constructor(options) {
@@ -59,19 +60,23 @@ class CodeRunner {
     this.runCode.registerVar(varName, code);
   }
 
-  evalCode(code, cb, forConsoleOnly = false) {
+  async evalCode(code, cb, forConsoleOnly = false) {
     cb = cb || function() {};
     const awaitIdx = code.indexOf('await');
+    let awaiting = false;
+
     if (awaitIdx > -1) {
-      if (awaitIdx < 2) {
-        let end = code.length;
-        if (code[end - 1] === ';') {
-          end--; // Remove the `;` because we add function calls
-        }
-        code = code.substring(5, end); // remove await keyword
+      awaiting = true;
+      const instructions = Utils.compact(code.split(';'));
+      const last = instructions.pop();
+
+      if (!last.trim().startsWith('return')) {
+        instructions.push(`return ${last}`); 
       } else {
-        code = `(async function() {${code}})();`;
+        instructions.push(last); 
       }
+      
+      code = `(async function() {${instructions.join(';')}})();`;
     }
     let result = this.runCode.doEval(code);
 
@@ -80,11 +85,16 @@ class CodeRunner {
       this.ipc.broadcast("runcode:newCommand", {code});
     }
 
-    if (result instanceof Promise && !result.on) {
-      return result.then((value) => { cb(null, value); }).catch(cb);
+    if (!awaiting) {
+      return cb(null, result);
     }
 
-    cb(null, result);
+    try {
+      const value = await result;
+      cb(null, value);
+    } catch (error) {
+      cb(error);  
+    }
   }
 }
 

From b2206b9358b79f365a8ad6895c97d28b2fae3de5 Mon Sep 17 00:00:00 2001
From: Anthony Laibe <anthony@laibe.cc>
Date: Mon, 17 Sep 2018 10:13:40 +0100
Subject: [PATCH 2/2] Fix test

---
 test_apps/test_app/config/contracts.js | 2 +-
 test_apps/test_app/test/token_spec.js  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/test_apps/test_app/config/contracts.js b/test_apps/test_app/config/contracts.js
index 8b4a36bcd..c68c002ef 100644
--- a/test_apps/test_app/config/contracts.js
+++ b/test_apps/test_app/config/contracts.js
@@ -53,7 +53,7 @@ module.exports = {
         }
       },
       SomeContract: {
-        deployIf: 'MyToken.methods.isAvailable().call()',
+        deployIf: 'await MyToken.methods.isAvailable().call()',
         args: [
           ["$MyToken2", "$SimpleStorage"],
           100
diff --git a/test_apps/test_app/test/token_spec.js b/test_apps/test_app/test/token_spec.js
index 8d4dd32e1..6dbdcd1f9 100644
--- a/test_apps/test_app/test/token_spec.js
+++ b/test_apps/test_app/test/token_spec.js
@@ -41,7 +41,7 @@ config({
       }
     },
     SomeContract: {
-      deployIf: "MyToken.methods.isAvailable().call()",
+      deployIf: "await MyToken.methods.isAvailable().call()",
       args: [
         ["$MyToken2", "$SimpleStorage"],
         100