From c3e068c4039f5163aeed250d44914664855fd6d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89ric=20Lemoine?= Date: Fri, 24 Jan 2020 16:29:41 +0100 Subject: [PATCH] [WIP] Fix bug in get_connection --- procrastinate/aiopg_connector.py | 2 +- tests/integration/test_pg_store.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/procrastinate/aiopg_connector.py b/procrastinate/aiopg_connector.py index da6e54b49..c19d77a40 100644 --- a/procrastinate/aiopg_connector.py +++ b/procrastinate/aiopg_connector.py @@ -103,7 +103,7 @@ def __init__(self, *, socket_timeout: float = store.SOCKET_TIMEOUT, **kwargs: An self.socket_timeout = socket_timeout async def get_connection(self, loads=None): - if not self._connection: + if not self._connection or self._connection.closed: self._connection = await get_connection( loads=loads, **self._connection_parameters ) diff --git a/tests/integration/test_pg_store.py b/tests/integration/test_pg_store.py index 99e14711a..bd248ac19 100644 --- a/tests/integration/test_pg_store.py +++ b/tests/integration/test_pg_store.py @@ -525,3 +525,16 @@ async def test_close_connection_no_connection(pg_job_store): async def test_stop_no_connection(pg_job_store): pg_job_store.stop() # Well we didn't crash. Great. + + +async def test_get_and_close_connection(pg_job_store): + conn1 = await pg_job_store.get_connection(loads=None) + assert not conn1.closed + conn2 = await pg_job_store.get_connection(loads=None) + assert conn2 is conn1 + await pg_job_store.close_connection() + assert conn1.closed + conn2 = await pg_job_store.get_connection(loads=None) + assert conn2 is not conn1 + await pg_job_store.close_connection() + assert conn2.closed