1 files changed, 21 insertions, 17 deletions
diff --git a/tests/integration/base.py b/tests/integration/base.py
index 3fb25b5..ea43d05 100644
--- a/tests/integration/base.py
+++ b/tests/integration/base.py
@@ -2,6 +2,7 @@ import shutil
 import unittest
 
 import docker
+from docker.utils import kwargs_from_env
 import six
 
 
@@ -10,20 +11,14 @@ BUSYBOX = 'busybox:buildroot-2014.02'
 
 class BaseIntegrationTest(unittest.TestCase):
     """
-    A base class for integration test cases.
-
-    It sets up a Docker client and cleans up the Docker server after itself.
+    A base class for integration test cases. It cleans up the Docker server
+    after itself.
     """
-    tmp_imgs = []
-    tmp_containers = []
-    tmp_folders = []
-    tmp_volumes = []
 
     def setUp(self):
         if six.PY2:
             self.assertRegex = self.assertRegexpMatches
             self.assertCountEqual = self.assertItemsEqual
-        self.client = docker.from_env(timeout=60)
         self.tmp_imgs = []
         self.tmp_containers = []
         self.tmp_folders = []
@@ -31,32 +26,41 @@ class BaseIntegrationTest(unittest.TestCase):
         self.tmp_networks = []
 
     def tearDown(self):
+        client = docker.from_env()
         for img in self.tmp_imgs:
             try:
-                self.client.remove_image(img)
+                client.api.remove_image(img)
             except docker.errors.APIError:
                 pass
         for container in self.tmp_containers:
             try:
-                self.client.stop(container, timeout=1)
-                self.client.remove_container(container)
+                client.api.remove_container(container, force=True)
             except docker.errors.APIError:
                 pass
         for network in self.tmp_networks:
             try:
-                self.client.remove_network(network)
+                client.api.remove_network(network)
             except docker.errors.APIError:
                 pass
-        for folder in self.tmp_folders:
-            shutil.rmtree(folder)
-
         for volume in self.tmp_volumes:
             try:
-                self.client.remove_volume(volume)
+                client.api.remove_volume(volume)
             except docker.errors.APIError:
                 pass
 
-        self.client.close()
+        for folder in self.tmp_folders:
+            shutil.rmtree(folder)
+
+
+class BaseAPIIntegrationTest(BaseIntegrationTest):
+    """
+    A test case for `APIClient` integration tests. It sets up an `APIClient`
+    as `self.client`.
+    """
+
+    def setUp(self):
+        super(BaseAPIIntegrationTest, self).setUp()
+        self.client = docker.APIClient(timeout=60, **kwargs_from_env())
 
     def run_container(self, *args, **kwargs):
         container = self.client.create_container(*args, **kwargs)