diff --git a/README.md b/README.md
index 454f44a..027fb6d 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,10 @@
 # YouTube Transcript/Subtitle API (including automatically generated subtitles)
 
+[![Build Status](https://travis-ci.org/jdepoix/youtube-transcript-api.svg)](https://travis-ci.org/jdepoix/youtube-transcript-api)
+[![MIT license](http://img.shields.io/badge/license-MIT-brightgreen.svg?style=flat)
+[![image](https://img.shields.io/pypi/v/youtube-transcript-api.svg)](https://pypi.org/project/youtube-transcript-api/)
+[![image](https://img.shields.io/pypi/pyversions/youtube-transcript-api.svg)](https://pypi.org/project/youtube-transcript-api/)
+
 This is an python API which allows you to get the transcripts/subtitles for a given YouTube video. It also works for automatically generated subtitles and it does not require a headless browser, like other selenium based solutions do!
 
 ## Install
diff --git a/requirements.txt b/requirements.txt
index a46f363..394e537 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 requests
 
 # testing
+mock
 httpretty
 coverage
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 9fd8506..e5895e0 100644
--- a/setup.py
+++ b/setup.py
@@ -43,6 +43,7 @@ setuptools.setup(
         'requests',
     ],
     tests_require=[
+        'mock',
         'httpretty',
         'coverage',
     ],
diff --git a/youtube_transcript_api/test/test_api.py b/youtube_transcript_api/test/test_api.py
index 76e9d19..a8105e7 100644
--- a/youtube_transcript_api/test/test_api.py
+++ b/youtube_transcript_api/test/test_api.py
@@ -1,5 +1,5 @@
 from unittest import TestCase
-from unittest.mock import MagicMock
+from mock import MagicMock
 
 import os
 
diff --git a/youtube_transcript_api/test/test_cli.py b/youtube_transcript_api/test/test_cli.py
index 72c5890..8a63c23 100644
--- a/youtube_transcript_api/test/test_cli.py
+++ b/youtube_transcript_api/test/test_cli.py
@@ -1,5 +1,5 @@
 from unittest import TestCase
-from unittest.mock import MagicMock
+from mock import MagicMock
 
 import json