Make tests a little more robust

kyuupichan · Apr 30, 2019 · a7f953e · a7f953e
1 parent 142bcd9
commit a7f953e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/aiorpcx/session.py b/aiorpcx/session.py
@@ -261,12 +261,12 @@ def connection_lost(self, exc):
         # Work around uvloop bug; see https://github.com/MagicStack/uvloop/issues/246
         if self.transport:
             self.transport = None
-            self.closed_event.set()
             # Release waiting tasks
             self._can_send.set()
             # Cancelling directly leads to self-cancellation problems for member
             # functions await-ing self.close()
             self.loop.call_soon(self._task.cancel)
+            self.closed_event.set()
 
     # External API
     def is_send_buffer_full(self):
@@ -349,6 +349,7 @@ async def close(self, *, force_after=30):
             except TaskTimeout:
                 self.abort()
                 await self.closed_event.wait()
+            await sleep(0)
 
 
 class MessageSession(SessionBase):

diff --git a/tests/test_session.py b/tests/test_session.py
@@ -598,9 +598,6 @@ async def test_log_me(self, server, caplog):
 
 class RPCClient(RPCSession):
     # For tests of wire messages
-    def connection_made(self, transport):
-        self.transport = transport
-
     async def send(self, item):
         if not isinstance(item, str):
             item = json.dumps(item)

diff --git a/tests/test_socks.py b/tests/test_socks.py
@@ -552,8 +552,9 @@ async def test_good_SOCKS4(self, proxy_address, auth):
 
     @pytest.mark.asyncio
     async def test_auto_detect_at_address_failure(self):
-        result = await SOCKSProxy.auto_detect_at_address('8.8.8.8:53', None)
-        assert result is None
+        async with ignore_after(1.0):
+            result = await SOCKSProxy.auto_detect_at_address('8.8.8.8:53', None)
+            assert result is None
 
     @pytest.mark.asyncio
     async def test_auto_detect_at_address_cannot_connect(self):