diff --git a/index.d.ts b/index.d.ts
index 24cf875..76ed55e 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -1,18 +1,20 @@
-import {Middleware, Dispatch} from "redux";
+import { Middleware, Action, AnyAction } from "redux";
 
-
-export type ThunkAction<R, S, E> = (dispatch: Dispatch<S>, getState: () => S,
-                                    extraArgument: E) => R;
-
-declare module "redux" {
-  export interface Dispatch<S> {
-    <R, E>(asyncAction: ThunkAction<R, S, E>): R;
-  }
+export interface ThunkDispatch<S, E, A extends Action> {
+  <T extends A>(action: T): T;
+  <R>(asyncAction: ThunkAction<R, S, E, A>): R;
 }
 
+export type ThunkAction<R, S, E, A extends Action> = (
+  dispatch: ThunkDispatch<S, E, A>,
+  getState: () => S,
+  extraArgument: E
+) => R;
 
-declare const thunk: Middleware & {
-  withExtraArgument(extraArgument: any): Middleware;
-};
+export type ThunkMiddleware<S = {}, A extends Action = AnyAction, E = undefined> = Middleware<ThunkDispatch<S, E, A>, S, ThunkDispatch<S, E, A>>;
+
+declare const thunk: ThunkMiddleware & {
+  withExtraArgument<E>(extraArgument: E): ThunkMiddleware<{}, AnyAction, E>
+}
 
 export default thunk;
diff --git a/package.json b/package.json
index 2c75bd6..47a159f 100644
--- a/package.json
+++ b/package.json
@@ -69,10 +69,10 @@
     "eslint-config-airbnb": "1.0.2",
     "eslint-plugin-react": "^4.1.0",
     "mocha": "^2.2.5",
-    "redux": "^3.4.0",
+    "redux": "~4.0.0",
     "rimraf": "^2.5.2",
-    "typescript": "^1.8.10",
-    "typescript-definition-tester": "0.0.4",
+    "typescript": "~2.6.2",
+    "typings-tester": "^0.3.1",
     "webpack": "^1.12.14"
   }
 }
diff --git a/test/index.js b/test/index.js
index ce7c57d..10d1d49 100644
--- a/test/index.js
+++ b/test/index.js
@@ -1,6 +1,6 @@
 import chai from 'chai';
 import thunkMiddleware from '../src/index';
-import * as tt from 'typescript-definition-tester';
+import { checkDirectory } from 'typings-tester';
 
 
 describe('thunk middleware', () => {
@@ -97,12 +97,8 @@ describe('thunk middleware', () => {
   describe('TypeScript definitions', function test() {
     this.timeout(0);
 
-    it('should compile against index.d.ts', (done) => {
-      tt.compileDirectory(
-        __dirname,
-        fileName => fileName.match(/\.ts$/),
-        () => done()
-      );
+    it('should compile against index.d.ts', () => {
+      checkDirectory(__dirname);
     });
   });
 });
diff --git a/test/tsconfig.json b/test/tsconfig.json
new file mode 100644
index 0000000..61a6e65
--- /dev/null
+++ b/test/tsconfig.json
@@ -0,0 +1,8 @@
+{
+  "compilerOptions": {
+    "strict": true
+  },
+  "files": [
+    "./*.ts"
+  ]
+}
diff --git a/test/typescript.ts b/test/typescript.ts
index 816bf47..b263f50 100644
--- a/test/typescript.ts
+++ b/test/typescript.ts
@@ -1,33 +1,75 @@
-import {Store, Middleware} from 'redux';
-import thunk, {ThunkAction} from '../index.d.ts';
+import { createStore, applyMiddleware } from 'redux';
+import thunk, { ThunkAction, ThunkMiddleware } from '../index';
 
+type State = {
+  foo: string;
+};
 
-declare const store: Store<{foo: string}>;
+type Actions = { type: 'FOO' } | { type: 'BAR', result: number };
 
-store.dispatch(dispatch => {
-  dispatch({type: 'FOO'});
-});
+type ThunkResult<R> = ThunkAction<R, State, undefined, Actions>;
 
-store.dispatch((dispatch, getState) => {
-  const state = getState();
+const initialState: State = {
+  foo: 'foo'
+};
 
-  const foo: string = state.foo;
-});
+function fakeReducer(state: State = initialState, action: Actions): State {
+  return state;
+}
 
-const middleware: Middleware = thunk.withExtraArgument('bar');
+const store = createStore(fakeReducer, applyMiddleware(thunk as ThunkMiddleware<State, Actions>));
 
-store.dispatch((dispatch, getState, extraArg) => {
-  console.log(extraArg);
+store.dispatch(dispatch => {
+  dispatch({ type: 'FOO' });
+  // typings:expect-error
+  dispatch({ type: 'BAR' })
+  dispatch({ type: 'BAR', result: 5 })
+  // typings:expect-error
+  store.dispatch({ type: 'BAZ'});
 });
 
-const thunkAction: ThunkAction<void, {foo: string}, {bar: number}> =
-  (dispatch, getState, extraArg) => {
-    const foo: string = getState().foo;
-    const bar: number = extraArg.bar;
-
-    dispatch({type: 'FOO'});
+function testGetState(): ThunkResult<void> {
+  return (dispatch, getState) => {
+    const state = getState();
+    const foo: string = state.foo;
+    dispatch({ type: 'FOO' });
+    // typings:expect-error
+    dispatch({ type: 'BAR'});
+    dispatch({ type: 'BAR', result: 5 });
+    // typings:expect-error
+    dispatch({ type: 'BAZ'});
+    // Can dispatch another thunk action
+    dispatch(anotherThunkAction());
   };
+}
 
-const thunkActionDispatchOnly: ThunkAction<void, {}, {}> = dispatch => {
-  dispatch({type: 'FOO'});
-};
+function anotherThunkAction(): ThunkResult<string> {
+  return (dispatch, getState) => {
+    dispatch({ type: 'FOO' });
+    return 'hello';
+  }
+}
+
+store.dispatch({ type: 'FOO' });
+// typings:expect-error
+store.dispatch({ type: 'BAR' })
+store.dispatch({ type: 'BAR', result: 5 })
+// typings:expect-error
+store.dispatch({ type: 'BAZ'});
+store.dispatch(testGetState());
+
+const storeThunkArg = createStore(
+  fakeReducer,
+  applyMiddleware(thunk.withExtraArgument('bar') as ThunkMiddleware<State, Actions, string>)
+);
+
+storeThunkArg.dispatch((dispatch, getState, extraArg) => {
+  const bar: string = extraArg;
+  store.dispatch({ type: 'FOO' });
+  // typings:expect-error
+  store.dispatch({ type: 'BAR' })
+  store.dispatch({ type: 'BAR', result: 5 })
+  // typings:expect-error
+  store.dispatch({ type: 'BAZ'});
+  console.log(extraArg);
+});