Refactor created_client contextmanager

tests/test_client.py

@@ -6,10 +6,10 @@
 from clickhouse_driver.compression.zstd import Compressor as ZSTDCompressor
 from clickhouse_driver.protocol import Compression
 from tests.numpy.util import check_numpy
-from tests.testcase import BaseTestCase
+from tests.testcase import TestCase
 
 
-class ClientFromUrlTestCase(BaseTestCase):
+class ClientFromUrlTestCase(TestCase):
     def assertHostsEqual(self, client, another, msg=None):
         self.assertEqual(list(client.connection.hosts), another, msg=msg)
 
@@ -227,9 +227,3 @@ def test_settings_is_important(self):
     def test_use_numpy(self):
         c = Client.from_url('clickhouse://host?use_numpy=true')
         self.assertTrue(c.connection.context.client_settings['use_numpy'])
-
-    def test_context_manager(self):
-        with self.client as c:
-            c.execute('SELECT 1')
-            self.assertTrue(c.connection.connected)
-        self.assertFalse(c.connection.connected)

tests/test_connect.py

@@ -174,6 +174,12 @@ def test_remember_current_database(self):
             rv = client.execute('SELECT currentDatabase()')
             self.assertEqual(rv, [('system', )])
 
+    def test_context_manager(self):
+        with self.created_client() as c:
+            c.execute('SELECT 1')
+            self.assertTrue(c.connection.connected)
+        self.assertFalse(c.connection.connected)
+
 
 class FakeBufferedReader(BufferedReader):
     def __init__(self, inputs, bufsize=128):

tests/testcase.py

@@ -68,14 +68,8 @@ def _create_client(self, **kwargs):
         client_kwargs.update(kwargs)
         return Client(self.host, **client_kwargs)
 
-    @contextmanager
     def created_client(self, **kwargs):
-        client = self._create_client(**kwargs)
-
-        try:
-            yield client
-        finally:
-            client.disconnect()
+        return self._create_client(**kwargs)
 
     @classmethod
     def setUpClass(cls):