Skip to content

Commit 0a23096

Browse files
committed
Add tests for connection closure listeners
1 parent 3a637f1 commit 0a23096

File tree

4 files changed

+53
-2
lines changed

4 files changed

+53
-2
lines changed

asyncpg/_testbase/fuzzer.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,10 @@ def _close_connection(self, connection):
145145
if conn_task is not None:
146146
conn_task.cancel()
147147

148+
def close_all_connections(self):
149+
for conn in list(self.connections):
150+
self._close_connection(conn)
151+
148152

149153
class Connection:
150154
def __init__(self, client_sock, backend_sock, proxy):

asyncpg/connection.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,11 +185,16 @@ def add_close_listener(self, callback):
185185
:param callable callback:
186186
A callable receiving one argument:
187187
**connection**: a Connection the callback is registered with.
188+
189+
.. versionadded:: 0.21.0
188190
"""
189191
self._close_listeners.add(callback)
190192

191193
def remove_close_listener(self, callback):
192-
"""Remove a listening callback for the connection closing."""
194+
"""Remove a listening callback for the connection closing.
195+
196+
.. versionadded:: 0.21.0
197+
"""
193198
self._close_listeners.discard(callback)
194199

195200
def get_server_pid(self):

tests/test_adversity.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,4 @@ def kill_connectivity():
7070
await asyncio.gather(
7171
*workers, return_exceptions=True)
7272
finally:
73-
pool.terminate()
73+
pool.terminate()

tests/test_listeners.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@
66

77

88
import asyncio
9+
import os
10+
import platform
11+
import sys
12+
import unittest
913

1014
from asyncpg import _testbase as tb
1115
from asyncpg import exceptions
@@ -272,3 +276,41 @@ def listener1(*args):
272276
pass
273277

274278
con.add_log_listener(listener1)
279+
280+
281+
@unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
282+
@unittest.skipIf(
283+
platform.system() == 'Windows' and
284+
sys.version_info >= (3, 8),
285+
'not compatible with ProactorEventLoop which is default in Python 3.8')
286+
class TestConnectionCloseListener(tb.ProxiedClusterTestCase):
287+
288+
async def test_connection_close_callback_called_on_remote(self):
289+
290+
called = False
291+
292+
def close_cb(con):
293+
nonlocal called
294+
called = True
295+
296+
con = await self.connect()
297+
con.add_close_listener(close_cb)
298+
self.proxy.close_all_connections()
299+
try:
300+
await con.fetchval('SELECT 1')
301+
except Exception:
302+
pass
303+
self.assertTrue(called)
304+
305+
async def test_connection_close_callback_called_on_local(self):
306+
307+
called = False
308+
309+
def close_cb(con):
310+
nonlocal called
311+
called = True
312+
313+
con = await self.connect()
314+
con.add_close_listener(close_cb)
315+
await con.close()
316+
self.assertTrue(called)

0 commit comments

Comments
 (0)