plishing

Author: chillaq

splitio/push/sse.py

@@ -274,8 +274,4 @@ async def shutdown(self):
         self._shutdown_requested = True
         sock = self._session.connector._loop._ssock
         sock.shutdown(socket.SHUT_RDWR)
-        await self._conn.close()
-        for task in asyncio.Task.all_tasks():
-            if not task.done():
-                if task._coro.cr_code.co_name == 'connect_split_sse_client':
-                    task.cancel()
+        await self._conn.close()

tests/push/test_sse.py

@@ -128,6 +128,7 @@ def runner():
 class SSEClientAsyncTests(object):
     """SSEClient test cases."""
 
+#    @pytest.mark.asyncio
     async def test_sse_client_disconnects(self):
         """Test correct initialization. Client ends the connection."""
         server = SSEMockServer()
@@ -139,16 +140,18 @@ async def callback(event):
             events.append(event)
 
         client = SSEClientAsync(callback)
+
         async def connect_split_sse_client():
             await client.start('http://127.0.0.1:' + str(server.port()))
 
-        asyncio.gather(connect_split_sse_client())
+        self._client_task = asyncio.gather(connect_split_sse_client())
         server.publish({'id': '1'})
         server.publish({'id': '2', 'event': 'message', 'data': 'abc'})
         server.publish({'id': '3', 'event': 'message', 'data': 'def'})
         server.publish({'id': '4', 'event': 'message', 'data': 'ghi'})
         await asyncio.sleep(1)
         await client.shutdown()
+        self._client_task.cancel()
         await asyncio.sleep(1)
 
         assert events == [
@@ -212,7 +215,7 @@ async def runner():
             """SSE client runner thread."""
             await client.start('http://127.0.0.1:' + str(server.port()))
 
-        client_task = asyncio.get_event_loop().create_task(runner())
+        client_task = asyncio.gather(runner())
 
         server.publish({'id': '1'})
         server.publish({'id': '2', 'event': 'message', 'data': 'abc'})