5 files changed, 27 insertions, 49 deletions
diff --git a/test/auth.spec.js b/test/auth.spec.js
index effc3473cb..13afbd40ca 100644
--- a/test/auth.spec.js
+++ b/test/auth.spec.js
@@ -19,14 +19,11 @@ describe('authentication', () => {
     };
 
     const chain = new IpcChain(mockIpc);
-    chain
-      .require('authenticate')
-      .withInputValidation((secret) => {
-        expect(secret).to.equal(credentials.sharedSecret);
-      })
-      .done();
+    chain.require('authenticate').withInputValidation((secret) => {
+      expect(secret).to.equal(credentials.sharedSecret);
+    });
 
-    chain.require('connect').done();
+    chain.require('connect');
 
     chain.onSuccessOrFailure(done);
 
diff --git a/test/autologin.spec.js b/test/autologin.spec.js
index debacc6c60..a2633dac9d 100644
--- a/test/autologin.spec.js
+++ b/test/autologin.spec.js
@@ -11,17 +11,11 @@ describe('autologin', () => {
     const randomAccountToken = '12345';
 
     const chain = new IpcChain(mockIpc);
-    chain
-      .require('getAccount')
-      .withReturnValue(randomAccountToken)
-      .done();
+    chain.require('getAccount').withReturnValue(randomAccountToken);
 
-    chain
-      .require('getAccountData')
-      .withInputValidation((num) => {
-        expect(num).to.equal(randomAccountToken);
-      })
-      .done();
+    chain.require('getAccountData').withInputValidation((num) => {
+      expect(num).to.equal(randomAccountToken);
+    });
 
     chain.onSuccessOrFailure(done);
 
diff --git a/test/helpers/IpcChain.js b/test/helpers/IpcChain.js
index 44a1c7b0d7..cc96986cd3 100644
--- a/test/helpers/IpcChain.js
+++ b/test/helpers/IpcChain.js
@@ -18,14 +18,16 @@ export class IpcChain {
   }
 
   require<R>(ipcCall: string): StepBuilder<R> {
+    const builder = new StepBuilder(ipcCall);
     this._expectedCalls.push(ipcCall);
-    return new StepBuilder(ipcCall, this._addStep.bind(this));
+    this._addStep(builder);
+
+    return builder;
   }
 
   _addStep<R>(step: StepBuilder<R>) {
-    const me = this;
-    this._mockIpc[step.ipcCall] = function() {
-      return new Promise((r) => me._stepPromiseCallback(step, r, arguments));
+    this._mockIpc[step.ipcCall] = (...args: Array<mixed>) => {
+      return new Promise((r) => this._stepPromiseCallback(step, r, args));
     };
   }
 
@@ -80,16 +82,14 @@ export class IpcChain {
 
 class StepBuilder<R> {
   ipcCall: string;
-  inputValidation: ?() => void;
+  inputValidation: ?(...args: Array<mixed>) => void;
   returnValue: ?R;
-  _cb: (StepBuilder<R>) => void;
 
-  constructor(ipcCall: string, cb: (StepBuilder<R>) => void) {
+  constructor(ipcCall: string) {
     this.ipcCall = ipcCall;
-    this._cb = cb;
   }
 
-  withInputValidation(iv: () => void): this {
+  withInputValidation(iv: (...args: Array<mixed>) => void): this {
     this.inputValidation = iv;
     return this;
   }
@@ -98,8 +98,4 @@ class StepBuilder<R> {
     this.returnValue = rv;
     return this;
   }
-
-  done() {
-    this._cb(this);
-  }
 }
diff --git a/test/login.spec.js b/test/login.spec.js
index acb38db2eb..4e4dd5ab55 100644
--- a/test/login.spec.js
+++ b/test/login.spec.js
@@ -17,19 +17,13 @@ describe('Logging in', () => {
     const { store, mockIpc, backend } = setupBackendAndStore();
 
     const chain = new IpcChain(mockIpc);
-    chain
-      .require('getAccountData')
-      .withInputValidation((an) => {
-        expect(an).to.equal('123');
-      })
-      .done();
+    chain.require('getAccountData').withInputValidation((an) => {
+      expect(an).to.equal('123');
+    });
 
-    chain
-      .require('setAccount')
-      .withInputValidation((an) => {
-        expect(an).to.equal('123');
-      })
-      .done();
+    chain.require('setAccount').withInputValidation((an) => {
+      expect(an).to.equal('123');
+    });
 
     chain.onSuccessOrFailure(done);
 
diff --git a/test/logout.spec.js b/test/logout.spec.js
index c7ed45adb2..5037a834e9 100644
--- a/test/logout.spec.js
+++ b/test/logout.spec.js
@@ -16,13 +16,10 @@ describe('logging out', () => {
     const { mockIpc, backend } = setupBackendAndStore();
 
     const chain = new IpcChain(mockIpc);
-    chain
-      .require('setAccount')
-      .withInputValidation((num) => {
-        expect(num).to.be.null;
-      })
-      .done();
-    chain.require('disconnect').done();
+    chain.require('setAccount').withInputValidation((num) => {
+      expect(num).to.be.null;
+    });
+    chain.require('disconnect');
     chain.onSuccessOrFailure(done);
 
     backend.logout();