diff --git a/index.d.ts b/index.d.ts
index 8294714..dafcf61 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -1,20 +1,24 @@
-import { Middleware, Dispatch, Action, AnyAction } from "redux";
+import { Middleware, Action, AnyAction } from "redux";
 
 
 export type ThunkAction<R, S = {}, E = {}, A extends Action<any> = AnyAction> = (
-  dispatch: Dispatch<A>,
+  dispatch: ThunkDispatch<S, E, A>,
   getState: () => S,
   extraArgument: E
 ) => R;
 
-declare module "redux" {
-  export interface Dispatch<A extends Action<any> = AnyAction> {
-    <R, E>(asyncAction: ThunkAction<R, {}, E, A>): R;
-  }
+interface ThunkDispatch<S = {}, E = {}, A extends Action = AnyAction> {
+  <T extends A>(action: T): T;
 }
 
-declare const thunk: Middleware & {
-  withExtraArgument(extraArgument: any): Middleware;
+interface ThunkDispatch<S = {}, E = {}, A extends Action = AnyAction> {
+  <R>(asyncAction: ThunkAction<R, S, E, A>): R;
+}
+
+type ThunkMiddleware<E, S = {}> = Middleware<ThunkDispatch<S, E>, S, ThunkDispatch<S, E>>;
+
+declare const thunk: ThunkMiddleware<undefined> & {
+  withExtraArgument<E>(extraArgument: E): ThunkMiddleware<E>
 };
 
 export default thunk;
diff --git a/test/typescript.ts b/test/typescript.ts
index f43a3c6..390f070 100644
--- a/test/typescript.ts
+++ b/test/typescript.ts
@@ -1,16 +1,20 @@
-import {Store, Middleware} from 'redux';
-import thunk, {ThunkAction} from '../index';
+import { createStore, applyMiddleware } from 'redux';
+import thunk, { ThunkAction } from '../index';
 
 type State = {
-  foo: string
+  foo: string;
 };
 
 type Actions = { type: 'FOO' };
 
-declare const store: Store<State, Actions>;
+function fakeReducer(state: State, action: Actions): State {
+  return state;
+}
+
+const store = createStore(fakeReducer, applyMiddleware(thunk));
 
 store.dispatch(dispatch => {
-  dispatch({type: 'FOO'});
+  dispatch({ type: 'FOO' });
 });
 
 function testGetState(): ThunkAction<void, State, {}, Actions> {
@@ -22,20 +26,27 @@ function testGetState(): ThunkAction<void, State, {}, Actions> {
 
 store.dispatch(testGetState());
 
-const middleware: Middleware = thunk.withExtraArgument('bar');
+const storeThunkArg = createStore(
+  fakeReducer,
+  applyMiddleware(thunk.withExtraArgument('bar'))
+);
 
-store.dispatch((dispatch, getState, extraArg) => {
+storeThunkArg.dispatch((dispatch, getState, extraArg) => {
+  const bar: string = extraArg;
   console.log(extraArg);
 });
 
-const thunkAction: ThunkAction<void, {foo: string}, {bar: number}> =
-  (dispatch, getState, extraArg) => {
-    const foo: string = getState().foo;
-    const bar: number = extraArg.bar;
+const thunkAction: ThunkAction<void, State, { bar: number }, Actions> = (
+  dispatch,
+  getState,
+  extraArg
+) => {
+  const foo: string = getState().foo;
+  const bar: number = extraArg.bar;
 
-    dispatch({type: 'FOO'});
-  };
+  dispatch({ type: 'FOO' });
+};
 
-const thunkActionDispatchOnly: ThunkAction<void, {}, {}> = dispatch => {
-  dispatch({type: 'FOO'});
+const thunkActionDispatchOnly: ThunkAction<void, {}, {}, Actions> = dispatch => {
+  dispatch({ type: 'FOO' });
 };