WIP

Author: NoahStapp

test/asynchronous/test_async_cancellation.py

@@ -17,32 +17,55 @@
 
 import asyncio
 import sys
-from test.utils import async_get_pool, get_pool, one
+import traceback
+
+from test.utils import async_get_pool, get_pool, one, delay
 
 sys.path[0:0] = [""]
 
-from test.asynchronous import AsyncIntegrationTest, connected
+from test.asynchronous import AsyncIntegrationTest, connected, async_client_context
 
 
 class TestAsyncCancellation(AsyncIntegrationTest):
-    async def test_async_cancellation_does_not_close_connection(self):
-        client = await self.async_rs_or_single_client(maxPoolSize=1, retryReads=False)
+    async def test_async_cancellation_closes_connection(self):
+        client = await self.async_rs_or_single_client()
         pool = await async_get_pool(client)
         await connected(client)
         conn = one(pool.conns)
 
         async def task():
-            while True:
-                await client.db.test.insert_one({"x": 1})
-                await asyncio.sleep(0.005)
+            await client.db.test.find_one({"$where": delay(1.0)})
+
+        task = asyncio.create_task(task())
+
+        await asyncio.sleep(0.1)
+
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await task
+
+        self.assertTrue(conn.closed)
+
+    @async_client_context.require_transactions
+    async def test_async_cancellation_aborts_transaction(self):
+        client = await self.async_rs_or_single_client()
+        await connected(client)
+
+        session = client.start_session()
+
+        async def callback(session):
+            await client.db.test.find_one({"$where": delay(1.0)})
+
+        async def task():
+            await session.with_transaction(callback)
 
         task = asyncio.create_task(task())
 
-        # Make sure the task successfully runs a few operations to simulate a long-running user task
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.1)
 
         task.cancel()
         with self.assertRaises(asyncio.CancelledError):
             await task
 
-        self.assertFalse(conn.closed)
+        self.assertFalse(session.in_transaction)
+