Commit f7649296 authored by Joshua Roesslein's avatar Joshua Roesslein
Browse files

Merge pull request #658 from wjt/streaming-ReadBuffer-py3-pr-635

Test case for PR #635 (and bug #615) 
parents 8346123f 3a061c4b
......@@ -120,13 +120,13 @@ class TweepyStreamTests(unittest.TestCase):
self.assertEqual(u'Caf\xe9'.encode('utf8'), s.session.params['follow'])
class TweepyStreamReadBuffer(unittest.TestCase):
class TweepyStreamReadBufferTests(unittest.TestCase):
stream = """11\n{id:12345}\n\n24\n{id:23456, test:"blah"}\n"""
stream = six.b("""11\n{id:12345}\n\n24\n{id:23456, test:"blah"}\n""")
def test_read_tweet(self):
for length in [1, 2, 5, 10, 20, 50]:
buf = ReadBuffer(six.StringIO(self.stream), length)
buf = ReadBuffer(six.BytesIO(self.stream), length)
self.assertEqual('11\n', buf.read_line())
self.assertEqual('{id:12345}\n', buf.read_len(11))
self.assertEqual('\n', buf.read_line())
......@@ -157,13 +157,14 @@ class TweepyStreamReadBuffer(unittest.TestCase):
return ""
# Create a fake stream
stream = six.StringIO('')
stream = six.BytesIO(six.b(''))
# 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):
stream.close()
with patch.multiple(stream, create=True, read=mock_read):
# 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)
......@@ -175,14 +176,14 @@ class TweepyStreamReadBuffer(unittest.TestCase):
self.assertEqual(mock_read.call_count, 0)
def test_read_unicode_tweet(self):
stream = '11\n{id:12345}\n\n23\n{id:23456, test:"\xe3\x81\x93"}\n\n'
stream = six.b('11\n{id:12345}\n\n23\n{id:23456, test:"\xe3\x81\x93"}\n\n')
for length in [1, 2, 5, 10, 20, 50]:
buf = ReadBuffer(six.StringIO(stream), length)
buf = ReadBuffer(six.BytesIO(stream), length)
self.assertEqual('11\n', buf.read_line())
self.assertEqual('{id:12345}\n', buf.read_len(11))
self.assertEqual('\n', buf.read_line())
self.assertEqual('23\n', buf.read_line())
self.assertEqual('{id:23456, test:"\xe3\x81\x93"}\n', buf.read_len(23))
self.assertEqual(u'{id:23456, test:"\u3053"}\n', buf.read_len(23))
class TweepyStreamBackoffTests(unittest.TestCase):
......
......@@ -7,6 +7,7 @@
from __future__ import absolute_import, print_function
import logging
import re
import requests
from requests.exceptions import Timeout
from threading import Thread
......@@ -148,10 +149,11 @@ class ReadBuffer(object):
use small chunks so it can read the length and the tweet in 2 read calls.
"""
def __init__(self, stream, chunk_size):
def __init__(self, stream, chunk_size, encoding='utf-8'):
self._stream = stream
self._buffer = ''
self._buffer = six.b('')
self._chunk_size = chunk_size
self._encoding = encoding
def read_len(self, length):
while not self._stream.closed:
......@@ -160,7 +162,13 @@ class ReadBuffer(object):
read_len = max(self._chunk_size, length - len(self._buffer))
self._buffer += self._stream.read(read_len)
def read_line(self, sep='\n'):
def read_line(self, sep=six.b('\n')):
"""Read the data stream until a given separator is found (default \n)
:param sep: Separator to read until. Must by of the bytes type (str in python 2,
bytes in python 3)
:return: The str of the data read until sep
"""
start = 0
while not self._stream.closed:
loc = self._buffer.find(sep, start)
......@@ -173,7 +181,7 @@ class ReadBuffer(object):
def _pop(self, length):
r = self._buffer[:length]
self._buffer = self._buffer[length:]
return r
return r.decode(self._encoding)
class Stream(object):
......@@ -290,7 +298,14 @@ class Stream(object):
self.running = False
def _read_loop(self, resp):
buf = ReadBuffer(resp.raw, self.chunk_size)
charset = resp.headers.get('content-type', default='')
enc_search = re.search('charset=(?P<enc>\S*)', charset)
if enc_search is not None:
encoding = enc_search.group('enc')
else:
encoding = 'utf-8'
buf = ReadBuffer(resp.raw, self.chunk_size, encoding=encoding)
while self.running and not resp.raw.closed:
length = 0
......
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