Commit 485f64cc authored by Joshua Roesslein's avatar Joshua Roesslein
Browse files

Merge pull request #561 from michaelbrooks/stream-disconnection

Prevent infinite read loop in tweepy Stream
parents f8fa8ec8 74bdefd4
......@@ -133,6 +133,47 @@ class TweepyStreamReadBuffer(unittest.TestCase):
self.assertEqual('24\n', buf.read_line())
self.assertEqual('{id:23456, test:"blah"}\n', buf.read_len(24))
def test_read_empty_buffer(self):
"""
Requests can be closed by twitter.
The ReadBuffer should not loop infinitely when this happens.
Instead it should return and let the outer _read_loop handle it.
"""
# If the test fails, we are in danger of an infinite loop
# so we need to do some work to block that from happening
class InfiniteLoopException(Exception):
pass
self.called_count = 0
call_limit = 5
def on_read(chunk_size):
self.called_count += 1
if self.called_count > call_limit:
# we have failed
raise InfiniteLoopException("Oops, read() was called a bunch of times")
return ""
# Create a fake stream
stream = six.StringIO('')
# Mock it's read function so it can't be called too many times
mock_read = MagicMock(side_effect=on_read)
try:
with patch.multiple(stream, create=True, read=mock_read, closed=True):
# Now the stream can't call 'read' more than call_limit times
# and it looks like a requests stream that is closed
buf = ReadBuffer(stream, 50)
buf.read_line("\n")
except InfiniteLoopException:
self.fail("ReadBuffer.read_line tried to loop infinitely.")
# The mocked function not have been called at all since the stream looks closed
self.assertEqual(mock_read.call_count, 0)
class TweepyStreamBackoffTests(unittest.TestCase):
def setUp(self):
......
......@@ -154,7 +154,7 @@ class ReadBuffer(object):
self._chunk_size = chunk_size
def read_len(self, length):
while True:
while not self._stream.closed:
if len(self._buffer) >= length:
return self._pop(length)
read_len = max(self._chunk_size, length - len(self._buffer))
......@@ -162,7 +162,7 @@ class ReadBuffer(object):
def read_line(self, sep='\n'):
start = 0
while True:
while not self._stream.closed:
loc = self._buffer.find(sep, start)
if loc >= 0:
return self._pop(loc + len(sep))
......@@ -292,9 +292,9 @@ class Stream(object):
def _read_loop(self, resp):
buf = ReadBuffer(resp.raw, self.chunk_size)
while self.running:
while self.running and not resp.raw.closed:
length = 0
while True:
while not resp.raw.closed:
line = buf.read_line().strip()
if not line:
self.listener.keep_alive() # keep-alive new lines are expected
......@@ -334,7 +334,7 @@ class Stream(object):
# self._data(next_status_obj.decode('utf-8'))
if resp.raw._fp.isclosed():
if resp.raw.closed:
self.on_closed(resp)
def _start(self, async):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment