From f56b190d2c51e16ea60df2352ebf85379cbb4a3c Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 5 Jan 2022 03:24:52 +0100
Subject: [PATCH] Raise PoolClosed on wait() on a closed pool

Previously it might have raised an assert, if, for instance, wait would
have failed and retried.
---
 docs/news_pool.rst                      |  1 +
 psycopg_pool/psycopg_pool/pool.py       |  2 ++
 psycopg_pool/psycopg_pool/pool_async.py |  2 ++
 tests/pool/test_pool.py                 | 11 +++++++++--
 tests/pool/test_pool_async.py           | 11 +++++++++--
 5 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index c7ad91298..dd055345e 100644
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -22,6 +22,7 @@ psycopg_pool 3.0.3 (unreleased)
 
 - Throw `!ValueError` if the pool `!min_size` is set to 0 (instead of
   hanging).
+- Throw `!PoolClosed` calling `~ConnectionPool.wait()` on a closed pool.
 
 
 Current release
diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 895067c56..7e0b88bce 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -78,6 +78,8 @@ class ConnectionPool(BasePool[Connection[Any]]):
         program to terminate in case the environment is not configured
         properly, rather than trying to stay up the hardest it can.
         """
+        self._check_open_getconn()
+
         with self._lock:
             assert not self._pool_full_event
             if len(self._pool) >= self._nconns:
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index 297505e05..e92209733 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -68,6 +68,8 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
             self._open()
 
     async def wait(self, timeout: float = 30.0) -> None:
+        self._check_open_getconn()
+
         async with self._lock:
             assert not self._pool_full_event
             if len(self._pool) >= self._nconns:
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index f4d249e7a..1a57367bf 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -8,7 +8,6 @@ from typing import Any, List, Tuple
 import pytest
 
 import psycopg
-from psycopg.errors import OperationalError
 from psycopg.pq import TransactionStatus
 from psycopg._compat import Counter
 
@@ -145,6 +144,14 @@ def test_wait_ready(dsn, monkeypatch):
         p.wait(0.0001)  # idempotent
 
 
+def test_wait_closed(dsn):
+    with pool.ConnectionPool(dsn) as p:
+        pass
+
+    with pytest.raises(pool.PoolClosed):
+        p.wait()
+
+
 @pytest.mark.slow
 def test_setup_no_timeout(dsn, proxy):
     with pytest.raises(pool.PoolTimeout):
@@ -749,7 +756,7 @@ def test_reopen(dsn):
     assert p._sched_runner is None
     assert not p._workers
 
-    with pytest.raises(OperationalError, match="cannot be reused"):
+    with pytest.raises(psycopg.OperationalError, match="cannot be reused"):
         p.open()
 
 
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 9a0a6d5a4..4d37bef77 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -7,7 +7,6 @@ from typing import Any, List, Tuple
 import pytest
 
 import psycopg
-from psycopg.errors import OperationalError
 from psycopg.pq import TransactionStatus
 from psycopg._compat import create_task, Counter
 
@@ -146,6 +145,14 @@ async def test_wait_ready(dsn, monkeypatch):
         await p.wait(0.0001)  # idempotent
 
 
+async def test_wait_closed(dsn):
+    async with pool.AsyncConnectionPool(dsn) as p:
+        pass
+
+    with pytest.raises(pool.PoolClosed):
+        await p.wait()
+
+
 @pytest.mark.slow
 async def test_setup_no_timeout(dsn, proxy):
     with pytest.raises(pool.PoolTimeout):
@@ -725,7 +732,7 @@ async def test_reopen(dsn):
     await p.close()
     assert p._sched_runner is None
 
-    with pytest.raises(OperationalError, match="cannot be reused"):
+    with pytest.raises(psycopg.OperationalError, match="cannot be reused"):
         await p.open()
 
 
-- 
2.47.2