diff --git a/.gitignore b/.gitignore
index f66473f..64a9605 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 *.xml
 *.swp
 *.tox
+.cache
 dist
 *egg-info
 target
diff --git a/docker_squash/lib/common.py b/docker_squash/lib/common.py
index 164e18a..f666182 100644
--- a/docker_squash/lib/common.py
+++ b/docker_squash/lib/common.py
@@ -6,6 +6,13 @@ import requests
 
 from docker_squash.errors import Error
 
+# First try to import Docker client using the API
+# available in version 2 of the library and fall
+# back to version 1
+try:
+    from docker.api.client import APIClient as APIClientClass
+except ImportError:
+    from docker.client import Client as APIClientClass
 
 DEFAULT_TIMEOUT_SECONDS = 600
 
@@ -31,10 +38,12 @@ def docker_client(log):
     except KeyError:
         pass
 
-    params = docker.utils.kwargs_from_env()
+    params = {"version": "auto"}
+    params.update(docker.utils.kwargs_from_env())
     params["timeout"] = timeout
+
     try:
-        client = docker.AutoVersionClient(**params)
+        client = APIClientClass(**params)
     except docker.errors.DockerException as e:
         log.error(
             "Could not create Docker client, please make sure that you specified valid parameters in the 'DOCKER_HOST' environment variable.")
diff --git a/requirements.txt b/requirements.txt
index a0dcd54..de5bcb3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,2 @@
-docker-py>=1.7.2
+docker
 six
diff --git a/tests/test_integ_squash.py b/tests/test_integ_squash.py
index aee5dce..bcab736 100644
--- a/tests/test_integ_squash.py
+++ b/tests/test_integ_squash.py
@@ -1,14 +1,11 @@
 import unittest
-import pytest
 import mock
 import six
 import codecs
-import docker
 import os
 import json
 import logging
 import shutil
-import sys
 import tarfile
 import io
 from io import BytesIO
@@ -16,6 +13,7 @@ import uuid
 
 from docker_squash.squash import Squash
 from docker_squash.errors import SquashError
+from docker_squash.lib import common
 
 if not six.PY3:
     import docker_squash.lib.xtarfile
@@ -40,10 +38,6 @@ class IntegSquash(unittest.TestCase):
 
     BUSYBOX_IMAGE = "busybox:1.24"
 
-    # Default base url for the connection
-    base_url = os.getenv('DOCKER_CONNECTION', 'unix://var/run/docker.sock')
-    docker = docker.AutoVersionClient(base_url=base_url)
-
     log = logging.getLogger()
     handler = logging.StreamHandler()
     formatter = logging.Formatter(
@@ -52,6 +46,8 @@ class IntegSquash(unittest.TestCase):
     log.addHandler(handler)
     log.setLevel(logging.DEBUG)
 
+    docker = common.docker_client(log)
+
     @classmethod
     def build_image(cls, dockerfile):
         IntegSquash.image = IntegSquash.Image(dockerfile)
@@ -529,7 +525,7 @@ class TestIntegSquash(IntegSquash):
             with self.SquashedImage(image, 2, output_path="image.tar"):
                 with tarfile.open("image.tar", mode='r') as tar:
                     squashed_layer_path = ImageHelper.top_layer_path(tar)
-                    
+
                     all_files = tar.getnames()
 
                     self.assertIn("%s/json" % squashed_layer_path, all_files)