diff --git a/package.json b/package.json
index 9ff44c9..ad79735 100644
--- a/package.json
+++ b/package.json
@@ -11,7 +11,7 @@
     "postinstall": "typings install",
     "build": "tsc",
     "test": "npm run build && npm run mocha && nsp check",
-    "mocha": "NODE_ENV=test mocha --compilers ts:ts-node/register --require scripts/test-helper.js src/__tests__/**.test.ts",
+    "mocha": "mocha --compilers ts:ts-node/register --require scripts/test-helper.js  src/__tests__/**.test.ts",
     "lint": "tslint src/**/*.ts"
   },
   "engines": {
@@ -43,6 +43,7 @@
   },
   "devDependencies": {
     "chai": "=3.5.0",
+    "chai-as-promised": "=6.0.0",
     "mocha": "=3.2.0",
     "node-fetch": "=1.6.3",
     "nsp": "=2.6.2",
diff --git a/scripts/test-helper.js b/scripts/test-helper.js
index 3856bdb..8cfab6a 100644
--- a/scripts/test-helper.js
+++ b/scripts/test-helper.js
@@ -1,3 +1,7 @@
 const chai = require('chai');
+const chaiAsPromised = require('chai-as-promised');
 
 chai.expect();
+chai.use(chaiAsPromised);
+
+process.env.NODE_ENV = 'test';
diff --git a/src/__tests__/server.test.ts b/src/__tests__/server.test.ts
index c92bccf..e53a9b5 100644
--- a/src/__tests__/server.test.ts
+++ b/src/__tests__/server.test.ts
@@ -33,7 +33,8 @@ describe('Server', function () {
            open: false
        }, '/foo/bar', function (response : any) {
            expect(response.ok).to.be.false;
-           done();
+           expect(response.status).to.equal(404);
+           expect(response.text()).to.eventually.include('Cannot GET').notify(done);
        });
     });
 
@@ -51,10 +52,7 @@ describe('Server', function () {
                     open: false
                 }, path, function (response : any) {
                     expect(response.status).to.equal(200);
-                    response.text().then(function (text : any) {
-                        expect(text).to.equal(body);
-                        done();
-                    });
+                    expect(response.text()).to.eventually.equal(body).notify(done);
                 });
             });
         });
diff --git a/src/middleware/404.ts b/src/middleware/404.ts
index 3bb2a29..4fe6100 100644
--- a/src/middleware/404.ts
+++ b/src/middleware/404.ts
@@ -1,6 +1,6 @@
 import { Request, Response } from 'express';
-import staticFile from 'connect-static-file'
-import path from 'path';
+import * as staticFile from 'connect-static-file'
+import * as path from 'path';
 
 export default function handle404(serveDir : string) {
     const handle404Middleware = staticFile(path.join(serveDir, '.404.html'));
diff --git a/src/server.ts b/src/server.ts
index cbb6801..6753f79 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -39,7 +39,7 @@ export default function createServer(opts : Options) : express.Application {
     }
 
     app.use(staticFileHandle(opts.serveDir));
-    app.use(handle404);
+    app.use(handle404(opts.serveDir));
 
     app.use(compression({ level: 9 }));
     app.use(helmet());
diff --git a/src/types/fakes.d.ts b/src/types/fakes.d.ts
index 810a1dd..4f0a349 100644
--- a/src/types/fakes.d.ts
+++ b/src/types/fakes.d.ts
@@ -9,3 +9,4 @@ declare module 'opbeat';
 declare module 'docopt';
 declare module 'open';
 declare module 'node-fetch';
+declare module 'chai';
diff --git a/tsconfig.json b/tsconfig.json
index a2a5e35..f2c86b3 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -17,7 +17,8 @@
   "exclude": [
     "dist",
     "node_modules",
-    "scripts/"
+    "scripts/",
+    "src/__tests__"
   ],
   "typeRoots": [
     "node_modules",