def test_civis_to_file_local(mock_requests):
# Test that a call to civis_to_file uses `requests` to grab the contents
# of a URL given by the API client and writes it to a file.
mock_civis = create_client_mock()
mock_requests.get.return_value.iter_content.return_value =\
(l.encode() for l in 'abcdef')
with TemporaryDirectory() as tdir:
fname = os.path.join(tdir, 'testfile')
_files.civis_to_file(137, fname, client=mock_civis)
with open(fname, 'rt') as _fin:
assert _fin.read() == 'abcdef'
mock_civis.files.get.assert_called_once_with(137)
mock_requests.get.assert_called_once_with(
mock_civis.files.get.return_value.file_url, stream=True)
def test_civis_to_file_retries(mock_requests):
mock_civis = create_client_mock()
# Mock the request iter_content so it fails partway the first time.
# Python 2.7 doesn't have the nonlocal keyword, so here's a little class to
# track whether it's the first call.
class UnreliableIterContent:
def __init__(self):
self.first_try = True
def mock_iter_content(self, _):
chunks = [l.encode() for l in 'abcdef']
for i, chunk in enumerate(chunks):
# Fail partway through on the first try.
if self.first_try and i == 3:
self.first_try = False
raise requests.ConnectionError()
yield chunk
mock_requests.get.return_value.iter_content = \
UnreliableIterContent().mock_iter_content
# Add some data to the buffer to test that we seek to the right place
# when retrying.
buf = io.BytesIO(b'0123')
buf.seek(4)
_files.civis_to_file(137, buf, client=mock_civis)
# Check that retries work and that the buffer position is reset.
# If we didn't seek when retrying, we'd get abcabcdef.
# If we seek'd to position 0, then we'd get abcdef.
buf.seek(0)
assert buf.read() == b'0123abcdef'
mock_civis.files.get.assert_called_once_with(137)
assert mock_requests.get.call_count == 2
mock_requests.get.assert_called_with(
mock_civis.files.get.return_value.file_url, stream=True)
def test_civis_to_file_retries(mock_requests):
mock_civis = create_client_mock()
first_try = True
# Mock the request iter_content so it fails partway the first time.
def mock_iter_content(_):
nonlocal first_try
chunks = [l.encode() for l in 'abcdef']
for i, chunk in enumerate(chunks):
# Fail partway through on the first try.
if first_try and i == 3:
first_try = False
raise requests.ConnectionError()
yield chunk
mock_requests.get.return_value.iter_content = mock_iter_content
# Add some data to the buffer to test that we seek to the right place
# when retrying.
buf = io.BytesIO(b'0123')
buf.seek(4)
_files.civis_to_file(137, buf, client=mock_civis)
# Check that retries work and that the buffer position is reset.
# If we didn't seek when retrying, we'd get abcabcdef.
# If we seek'd to position 0, then we'd get abcdef.
buf.seek(0)
assert buf.read() == b'0123abcdef'
mock_civis.files.get.assert_called_once_with(137)
assert mock_requests.get.call_count == 2
mock_requests.get.assert_called_with(
mock_civis.files.get.return_value.file_url, stream=True)