Fix StreamReader._wait_for_data()

Return if we got data or EOF before waiting for data.

2 files changed, 12 insertions, 7 deletions

diff --git a/tests/test_streams.py b/tests/test_streams.py
index 827c583..35f0100 100644
--- a/tests/test_streams.py
+++ b/tests/test_streams.py
@@ -18,9 +18,6 @@ from trollius import test_utils
 from trollius.test_utils import mock
 
 
-SHORT_DELAY = 0.010
-
-
 class StreamReaderTests(test_utils.TestCase):
 
     DATA = b'line1\nline2\nline3\n'
@@ -150,7 +147,7 @@ class StreamReaderTests(test_utils.TestCase):
 
         def cb():
             stream.feed_data(self.DATA)
-        self.loop.call_later(SHORT_DELAY, cb)
+        self.loop.call_soon(cb)
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(self.DATA, data)
@@ -174,7 +171,7 @@ class StreamReaderTests(test_utils.TestCase):
 
         def cb():
             stream.feed_eof()
-        self.loop.call_later(SHORT_DELAY, cb)
+        self.loop.call_soon(cb)
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(b'', data)
@@ -189,7 +186,7 @@ class StreamReaderTests(test_utils.TestCase):
             stream.feed_data(b'chunk1\n')
             stream.feed_data(b'chunk2')
             stream.feed_eof()
-        self.loop.call_later(SHORT_DELAY, cb)
+        self.loop.call_soon(cb)
 
         data = self.loop.run_until_complete(read_task)
 
@@ -218,7 +215,7 @@ class StreamReaderTests(test_utils.TestCase):
             stream.feed_data(b'chunk2 ')
             stream.feed_data(b'chunk3 ')
             stream.feed_data(b'\n chunk4')
-        self.loop.call_later(SHORT_DELAY, cb)
+        self.loop.call_soon(cb)
 
         line = self.loop.run_until_complete(read_task)
         self.assertEqual(b'chunk1 chunk2 chunk3 \n', line)
diff --git a/trollius/streams.py b/trollius/streams.py
index e83e567..b611c80 100644
--- a/trollius/streams.py
+++ b/trollius/streams.py
@@ -398,6 +398,7 @@ class StreamReader(object):
             else:
                 self._paused = True
 
+    @coroutine
     def _wait_for_data(self, func_name):
         """Wait until feed_data() or feed_eof() is called."""
         # StreamReader uses a future to link the protocol feed_data() method
@@ -408,6 +409,13 @@ class StreamReader(object):
             raise RuntimeError('%s() called while another coroutine is '
                                'already waiting for incoming data' % func_name)
 
+        # In asyncio, there is no need to recheck if we got data or EOF thanks
+        # to "yield from". In trollius, a StreamReader method can be called
+        # after the _wait_for_data() coroutine is scheduled and before it is
+        # really executed.
+        if self._buffer or self._eof:
+            return
+
         self._waiter = futures.Future(loop=self._loop)
         try:
             yield From(self._waiter)