Skip to content

Commit 3809c5b

Browse files
committed
tests: Make clusters easier to manage
The current `TestCase.start_cluster()` helper is not flexible enough as it does not allow making customizations to the cluster before it is started. The new `TestCase.new_cluster()` and `TestCase.start_cluster()` are now used instead. Cluster cleanup is now automatic.
1 parent c8f1e20 commit 3809c5b

File tree

4 files changed

+57
-36
lines changed

4 files changed

+57
-36
lines changed

asyncpg/_testbase/__init__.py

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -205,16 +205,21 @@ def _format_loop_exception(self, context, n):
205205
_default_cluster = None
206206

207207

208-
def _start_cluster(ClusterCls, cluster_kwargs, server_settings,
209-
initdb_options=None):
208+
def _init_cluster(ClusterCls, cluster_kwargs, initdb_options=None):
210209
cluster = ClusterCls(**cluster_kwargs)
211210
cluster.init(**(initdb_options or {}))
212211
cluster.trust_local_connections()
213-
cluster.start(port='dynamic', server_settings=server_settings)
214212
atexit.register(_shutdown_cluster, cluster)
215213
return cluster
216214

217215

216+
def _start_cluster(ClusterCls, cluster_kwargs, server_settings,
217+
initdb_options=None):
218+
cluster = _init_cluster(ClusterCls, cluster_kwargs, initdb_options)
219+
cluster.start(port='dynamic', server_settings=server_settings)
220+
return cluster
221+
222+
218223
def _get_initdb_options(initdb_options=None):
219224
if not initdb_options:
220225
initdb_options = {}
@@ -228,7 +233,7 @@ def _get_initdb_options(initdb_options=None):
228233
return initdb_options
229234

230235

231-
def _start_default_cluster(server_settings={}, initdb_options=None):
236+
def _init_default_cluster(initdb_options=None):
232237
global _default_cluster
233238

234239
if _default_cluster is None:
@@ -237,9 +242,8 @@ def _start_default_cluster(server_settings={}, initdb_options=None):
237242
# Using existing cluster, assuming it is initialized and running
238243
_default_cluster = pg_cluster.RunningCluster()
239244
else:
240-
_default_cluster = _start_cluster(
245+
_default_cluster = _init_cluster(
241246
pg_cluster.TempCluster, cluster_kwargs={},
242-
server_settings=server_settings,
243247
initdb_options=_get_initdb_options(initdb_options))
244248

245249
return _default_cluster
@@ -248,7 +252,8 @@ def _start_default_cluster(server_settings={}, initdb_options=None):
248252
def _shutdown_cluster(cluster):
249253
if cluster.get_status() == 'running':
250254
cluster.stop()
251-
cluster.destroy()
255+
if cluster.get_status() != 'not-initialized':
256+
cluster.destroy()
252257

253258

254259
def create_pool(dsn=None, *,
@@ -278,15 +283,40 @@ def get_server_settings(cls):
278283
'log_connections': 'on'
279284
}
280285

286+
@classmethod
287+
def new_cluster(cls, ClusterCls, *, cluster_kwargs={}, initdb_options={}):
288+
cluster = _init_cluster(ClusterCls, cluster_kwargs,
289+
_get_initdb_options(initdb_options))
290+
cls._clusters.append(cluster)
291+
return cluster
292+
293+
@classmethod
294+
def start_cluster(cls, cluster, *, server_settings={}):
295+
cluster.start(port='dynamic', server_settings=server_settings)
296+
281297
@classmethod
282298
def setup_cluster(cls):
283-
cls.cluster = _start_default_cluster(cls.get_server_settings())
299+
cls.cluster = _init_default_cluster()
300+
301+
if cls.cluster.get_status() != 'running':
302+
cls.cluster.start(
303+
port='dynamic', server_settings=cls.get_server_settings())
284304

285305
@classmethod
286306
def setUpClass(cls):
287307
super().setUpClass()
308+
cls._clusters = []
288309
cls.setup_cluster()
289310

311+
@classmethod
312+
def tearDownClass(cls):
313+
super().tearDownClass()
314+
for cluster in cls._clusters:
315+
if cluster is not _default_cluster:
316+
cluster.stop()
317+
cluster.destroy()
318+
cls._clusters = []
319+
290320
@classmethod
291321
def get_connection_spec(cls, kwargs={}):
292322
conn_spec = cls.cluster.get_connection_spec()
@@ -309,14 +339,6 @@ def connect(cls, **kwargs):
309339
conn_spec = cls.get_connection_spec(kwargs)
310340
return pg_connection.connect(**conn_spec, loop=cls.loop)
311341

312-
@classmethod
313-
def start_cluster(cls, ClusterCls, *,
314-
cluster_kwargs={}, server_settings={},
315-
initdb_options={}):
316-
return _start_cluster(
317-
ClusterCls, cluster_kwargs,
318-
server_settings, _get_initdb_options(initdb_options))
319-
320342

321343
class ProxiedClusterTestCase(ClusterTestCase):
322344
@classmethod

asyncpg/cluster.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ def get_status(self):
106106
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
107107
stdout, stderr = process.stdout, process.stderr
108108

109-
if process.returncode == 4 or not os.listdir(self._data_dir):
109+
if (process.returncode == 4 or not os.path.exists(self._data_dir) or
110+
not os.listdir(self._data_dir)):
110111
return 'not-initialized'
111112
elif process.returncode == 3:
112113
return 'stopped'

tests/test_connect.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -756,8 +756,9 @@ def get_server_settings(cls):
756756

757757
@classmethod
758758
def setup_cluster(cls):
759-
cls.cluster = cls.start_cluster(
760-
pg_cluster.TempCluster, server_settings=cls.get_server_settings())
759+
cls.cluster = cls.new_cluster(pg_cluster.TempCluster)
760+
cls.start_cluster(
761+
cls.cluster, server_settings=cls.get_server_settings())
761762

762763
def setUp(self):
763764
super().setUp()

tests/test_pool.py

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -743,17 +743,17 @@ class MyException(Exception):
743743

744744

745745
@unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
746-
class TestHotStandby(tb.ConnectedTestCase):
746+
class TestHotStandby(tb.ClusterTestCase):
747747
@classmethod
748-
def setUpClass(cls):
749-
super().setUpClass()
750-
751-
cls.master_cluster = cls.start_cluster(
752-
pg_cluster.TempCluster,
748+
def setup_cluster(cls):
749+
cls.master_cluster = cls.new_cluster(pg_cluster.TempCluster)
750+
cls.start_cluster(
751+
cls.master_cluster,
753752
server_settings={
754753
'max_wal_senders': 10,
755754
'wal_level': 'hot_standby'
756-
})
755+
}
756+
)
757757

758758
con = None
759759

@@ -771,27 +771,24 @@ def setUpClass(cls):
771771

772772
conn_spec = cls.master_cluster.get_connection_spec()
773773

774-
cls.standby_cluster = cls.start_cluster(
774+
cls.standby_cluster = cls.new_cluster(
775775
pg_cluster.HotStandbyCluster,
776776
cluster_kwargs={
777777
'master': conn_spec,
778778
'replication_user': 'replication'
779-
},
779+
}
780+
)
781+
cls.start_cluster(
782+
cls.standby_cluster,
780783
server_settings={
781784
'hot_standby': True
782-
})
785+
}
786+
)
783787

784788
finally:
785789
if con is not None:
786790
cls.loop.run_until_complete(con.close())
787791

788-
@classmethod
789-
def tearDownMethod(cls):
790-
cls.standby_cluster.stop()
791-
cls.standby_cluster.destroy()
792-
cls.master_cluster.stop()
793-
cls.master_cluster.destroy()
794-
795792
def create_pool(self, **kwargs):
796793
conn_spec = self.standby_cluster.get_connection_spec()
797794
conn_spec.update(kwargs)

0 commit comments

Comments
 (0)