Move the timeout handling into the buffer class

Author: dehort

receptor/connection/base.py

@@ -142,7 +142,7 @@ async def drain_buf(self, item):
 
     async def _wait_handshake(self):
         logger.debug("waiting for HI")
-        response = await asyncio.wait_for(self.buf.get(), 5.0)
+        response = await self.buf.get(timeout=5.0)
         self.remote_id = response.header["id"]
         await self.register()
         await self.receptor.recalculate_and_send_routes_soon()

receptor/messages/framed.py

@@ -287,8 +287,8 @@ async def finish(self):
         self.to_read = 0
         self.bb = FileBackedBuffer.from_temp()
 
-    async def get(self):
-        return await self.q.get()
+    async def get(self, timeout=None):
+        return await asyncio.wait_for(self.q.get(), timeout)
 
     def get_nowait(self):
         return self.q.get_nowait()