@@ -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+
218223def _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):
248252def _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
254259def 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
321343class ProxiedClusterTestCase (ClusterTestCase ):
322344 @classmethod
0 commit comments