Fix unit tests

Author: robsdedude

neo4j/_async_compat/shims/__init__.py

@@ -91,8 +91,8 @@ async def wait_for(fut, timeout):
                     # and https://github.com/python/cpython/pull/28149
 
                     # We got cancelled, but we are already done. Therefore,
-                    # we defer the cancellation until the task yields to the
-                    # event loop the next time.
+                    # we defer the cancellation until next time the task yields
+                    # to the event loop.
                     asyncio.current_task().cancel()
                     # [/PATCH]
                     return fut.result()

tests/integration/mixed/test_async_cancellation.py

@@ -99,21 +99,21 @@ async def _do_the_read(session_, i=1):
         raise
 
 
-REPEATS = 1000
+REPETITIONS = 1000
 
 
 @mark_async_test
 @pytest.mark.parametrize(("i", "read_func", "waits", "cancel_count"), (
     (
-        f"{i + 1:0{len(str(REPEATS))}}/{REPEATS}",
+        f"{i + 1:0{len(str(REPETITIONS))}}/{REPETITIONS}",
         random.choice((
             _do_the_read, _do_the_read_tx_context, _do_the_read_explicit_tx,
             _do_the_read_tx_func
         )),
         random.randint(0, 1000),
         random.randint(1, 20),
     )
-    for i in range(REPEATS)  # repeats
+    for i in range(REPETITIONS)
 ))
 async def test_async_cancellation(
     uri, auth, mocker, read_func, waits, cancel_count, i
@@ -175,7 +175,7 @@ async def test_async_cancellation(
             assert bookmarks != new_bookmarks
 
 
-SESSION_REPEATS = 100
+SESSION_REPETITIONS = 100
 READS_PER_SESSION = 20
 
 
@@ -187,7 +187,7 @@ async def test_async_cancellation_does_not_leak(uri, auth):
         # driver needs to cope with a single connection in the pool!
         max_connection_pool_size=1,
     ) as driver:
-        for session_number in range(SESSION_REPEATS):
+        for session_number in range(SESSION_REPETITIONS):
             async with driver.session() as session:
                 for read_number in range(READS_PER_SESSION):
                     read_func = random.choice((

tests/unit/async_/io/test_direct.py

@@ -52,6 +52,7 @@ class AsyncQuickConnection:
     def __init__(self, socket):
         self.socket = socket
         self.address = socket.getpeername()
+        self.local_port = self.address[1]
 
     @property
     def is_reset(self):

tests/unit/async_/work/test_transaction.py

@@ -40,7 +40,9 @@ async def test_transaction_context_when_committing(
 ):
     on_closed = mocker.AsyncMock()
     on_error = mocker.AsyncMock()
-    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error)
+    on_cancel = mocker.Mock()
+    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     async with tx as tx_:
@@ -70,7 +72,9 @@ async def test_transaction_context_with_explicit_rollback(
 ):
     on_closed = mocker.AsyncMock()
     on_error = mocker.AsyncMock()
-    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error)
+    on_cancel = mocker.Mock()
+    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     async with tx as tx_:
@@ -99,7 +103,9 @@ class OopsError(RuntimeError):
 
     on_closed = MagicMock()
     on_error = MagicMock()
-    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error)
+    on_cancel = MagicMock()
+    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     with pytest.raises(OopsError):
@@ -117,7 +123,9 @@ class OopsError(RuntimeError):
 async def test_transaction_run_takes_no_query_object(async_fake_connection):
     on_closed = MagicMock()
     on_error = MagicMock()
-    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error)
+    on_cancel = MagicMock()
+    tx = AsyncTransaction(async_fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     with pytest.raises(ValueError):
         await tx.run(Query("RETURN 1"))
 
@@ -128,7 +136,7 @@ async def test_transaction_rollbacks_on_open_connections(
 ):
     tx = AsyncTransaction(
         async_fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     async with tx as tx_:
         async_fake_connection.is_reset_mock.return_value = False
@@ -145,7 +153,7 @@ async def test_transaction_no_rollback_on_reset_connections(
 ):
     tx = AsyncTransaction(
         async_fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     async with tx as tx_:
         async_fake_connection.is_reset_mock.return_value = True
@@ -162,7 +170,7 @@ async def test_transaction_no_rollback_on_closed_connections(
 ):
     tx = AsyncTransaction(
         async_fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     async with tx as tx_:
         async_fake_connection.closed.return_value = True
@@ -181,7 +189,7 @@ async def test_transaction_no_rollback_on_defunct_connections(
 ):
     tx = AsyncTransaction(
         async_fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     async with tx as tx_:
         async_fake_connection.defunct.return_value = True

tests/unit/mixed/async_compat/test_concurrency.py

@@ -176,19 +176,13 @@ async def test_async_r_lock_acquire_cancellation(waits):
     async def acquire_task():
         while True:
             count = lock._count
-            cancellation = None
-            try:
-                print("try")
-                await lock.acquire(timeout=0.1)
-                print("acquired")
-            except asyncio.CancelledError as exc:
-                print("cancelled")
-                cancellation = exc
-
-            if cancellation is not None:
-                assert lock._count == count
-                raise cancellation
+            await lock.acquire(timeout=0.1)
             assert lock._count == count + 1
+            try:
+                await asyncio.sleep(0)
+            except asyncio.CancelledError:
+                raise
+            assert count < 50  # safety guard, we shouldn't ever get there!
 
     fut = asyncio.ensure_future(acquire_task())
     for _ in range(waits):

tests/unit/mixed/async_compat/test_shims.py

@@ -26,22 +26,33 @@
 from ...._async_compat import mark_async_test
 
 
+async def _check_wait_for(wait_for_, should_propagate_cancellation):
+    inner = asyncio.get_event_loop().create_future()
+    outer = wait_for_(inner, 0.1)
+    outer_future = asyncio.ensure_future(outer)
+    await asyncio.sleep(0)
+    inner.set_result(None)  # inner is done
+    outer_future.cancel()  # AND outer got cancelled
+
+    if should_propagate_cancellation:
+        with pytest.raises(asyncio.CancelledError):
+            await outer_future
+    else:
+        await outer_future
+
+
 @pytest.mark.skipif(
     sys.version_info < (3, 8),
     reason="wait_for is only broken in Python 3.8+"
 )
 @mark_async_test
 async def test_wait_for_shim_is_necessary_starting_from_3x8():
     # when this tests fails, the shim became superfluous
-    inner = asyncio.get_event_loop().create_future()
-    outer = asyncio.wait_for(inner, 0.1)
-    outer_future = asyncio.ensure_future(outer)
-    await asyncio.sleep(0)
-    inner.set_result(None)  # inner is done
-    outer_future.cancel()  # AND outer got cancelled
-
-    # this should propagate the cancellation, but it's broken :/
-    await outer_future
+    await _check_wait_for(
+        asyncio.wait_for,
+        # this should propagate the cancellation, but it's broken :/
+        should_propagate_cancellation=False
+    )
 
 
 @pytest.mark.skipif(
@@ -50,25 +61,16 @@ async def test_wait_for_shim_is_necessary_starting_from_3x8():
 )
 @mark_async_test
 async def test_wait_for_shim_is_not_necessary_prior_to_3x8():
-    inner = asyncio.get_event_loop().create_future()
-    outer = asyncio.wait_for(inner, 0.1)
-    outer_future = asyncio.ensure_future(outer)
-    await asyncio.sleep(0)
-    inner.set_result(None)  # inner is done
-    outer_future.cancel()  # AND outer got cancelled
-
-    with pytest.raises(asyncio.CancelledError):
-        await outer_future
+    await _check_wait_for(
+        asyncio.wait_for,
+        should_propagate_cancellation=True
+    )
 
 
 @mark_async_test
 async def test_wait_for_shim_propagates_cancellation():
-    inner = asyncio.get_event_loop().create_future()
-    outer = shims.wait_for(inner, 0.1)
-    outer_future = asyncio.ensure_future(outer)
-    await asyncio.sleep(0)
-    inner.set_result(None)  # inner is done
-    outer_future.cancel()  # AND outer got cancelled
-
-    with pytest.raises(asyncio.CancelledError):
-        await outer_future
+    # shim should always work regardless of the Python version
+    await _check_wait_for(
+        shims.wait_for,
+        should_propagate_cancellation=True
+    )

tests/unit/sync/io/test_direct.py

@@ -52,6 +52,7 @@ class QuickConnection:
     def __init__(self, socket):
         self.socket = socket
         self.address = socket.getpeername()
+        self.local_port = self.address[1]
 
     @property
     def is_reset(self):

tests/unit/sync/work/test_transaction.py

@@ -40,7 +40,9 @@ def test_transaction_context_when_committing(
 ):
     on_closed = mocker.Mock()
     on_error = mocker.Mock()
-    tx = Transaction(fake_connection, 2, on_closed, on_error)
+    on_cancel = mocker.Mock()
+    tx = Transaction(fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     with tx as tx_:
@@ -70,7 +72,9 @@ def test_transaction_context_with_explicit_rollback(
 ):
     on_closed = mocker.Mock()
     on_error = mocker.Mock()
-    tx = Transaction(fake_connection, 2, on_closed, on_error)
+    on_cancel = mocker.Mock()
+    tx = Transaction(fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     with tx as tx_:
@@ -99,7 +103,9 @@ class OopsError(RuntimeError):
 
     on_closed = MagicMock()
     on_error = MagicMock()
-    tx = Transaction(fake_connection, 2, on_closed, on_error)
+    on_cancel = MagicMock()
+    tx = Transaction(fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     mock_commit = mocker.patch.object(tx, "_commit", wraps=tx._commit)
     mock_rollback = mocker.patch.object(tx, "_rollback", wraps=tx._rollback)
     with pytest.raises(OopsError):
@@ -117,7 +123,9 @@ class OopsError(RuntimeError):
 def test_transaction_run_takes_no_query_object(fake_connection):
     on_closed = MagicMock()
     on_error = MagicMock()
-    tx = Transaction(fake_connection, 2, on_closed, on_error)
+    on_cancel = MagicMock()
+    tx = Transaction(fake_connection, 2, on_closed, on_error,
+                          on_cancel)
     with pytest.raises(ValueError):
         tx.run(Query("RETURN 1"))
 
@@ -128,7 +136,7 @@ def test_transaction_rollbacks_on_open_connections(
 ):
     tx = Transaction(
         fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     with tx as tx_:
         fake_connection.is_reset_mock.return_value = False
@@ -145,7 +153,7 @@ def test_transaction_no_rollback_on_reset_connections(
 ):
     tx = Transaction(
         fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     with tx as tx_:
         fake_connection.is_reset_mock.return_value = True
@@ -162,7 +170,7 @@ def test_transaction_no_rollback_on_closed_connections(
 ):
     tx = Transaction(
         fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     with tx as tx_:
         fake_connection.closed.return_value = True
@@ -181,7 +189,7 @@ def test_transaction_no_rollback_on_defunct_connections(
 ):
     tx = Transaction(
         fake_connection, 2, lambda *args, **kwargs: None,
-        lambda *args, **kwargs: None
+        lambda *args, **kwargs: None, lambda *args, **kwargs: None
     )
     with tx as tx_:
         fake_connection.defunct.return_value = True