def test_scp_error(self):
"""
If an SCP error occurs which is not a 'No space left on device' error, it must be raised.
"""
file_name = 'dataset.nc.tar.gz'
with mock.patch.object(utils.RemoteStorage, 'put') as mock_put:
mock_put.side_effect = scp.SCPException()
with mock.patch.object(utils.LocalStorage, 'get_file_size', return_value=1), \
mock.patch.object(utils.LocalStorage, 'get_block_size', return_value=4096), \
mock.patch.object(utils.RemoteStorage, 'free_space'), \
mock.patch.object(utils.RemoteStorage, '__init__', return_value=None), \
mock.patch.object(utils.RemoteStorage, '__del__', return_value=None):
with self.assertRaises(scp.SCPException):
tasks.publish((1, [file_name])) # pylint: disable=no-value-for-parameter
def generate():
buff_size = buffer_size
pos = 0
while pos < size:
# we have to make sure we don't read the final byte
if size - pos <= buff_size:
buff_size = size - pos
buf = channel.recv(buff_size)
pos += len(buf)
yield buf
msg = channel.recv(512)
channel.close()
if msg and msg[0:1] != b'\x00':
self._closeSCPClient()
raise scp.SCPException(scp.asunicode(msg[1:]))
def stream(self, remote_path, buffer_size=1024):
client = self._connectSCPClient()
channel = client._open()
channel.settimeout(client.socket_timeout)
channel.exec_command(
b"scp -f " + client.sanitize(scp.asbytes(remote_path))) # nosec
while not channel.closed:
# wait for command as long as we're open
channel.sendall('\x00')
msg = channel.recv(1024)
if not msg: # chan closed while recving
break
assert msg[-1:] == b'\n'
msg = msg[:-1]
code = msg[0:1]
# recv file
if code == b"C":
cmd = msg[1:]
parts = cmd.strip().split(b' ', 2)
size = int(parts[1])
channel.send(b'\x00')
try:
def generate():
buff_size = buffer_size
pos = 0
while pos < size:
# we have to make sure we don't read the final byte
if size - pos <= buff_size:
buff_size = size - pos
buf = channel.recv(buff_size)
pos += len(buf)
yield buf
msg = channel.recv(512)
channel.close()
if msg and msg[0:1] != b'\x00':
self._closeSCPClient()
raise scp.SCPException(scp.asunicode(msg[1:]))
return generate()
except SocketTimeout:
channel.close()
self._closeSCPClient()
raise scp.SCPException('Error receiving, socket.timeout')
def test_no_space_left_error(self):
"""
If a 'No space left on device error' occurs, the partially downloaded file must be removed
and the task must be retried.
"""
file_name = 'dataset.nc.tar.gz'
with mock.patch.object(utils.RemoteStorage, 'put') as mock_put, \
mock.patch.object(utils.RemoteStorage, 'remove') as mock_remove:
mock_put.side_effect = scp.SCPException('No space left on device')
with mock.patch.object(utils.LocalStorage, 'get_file_size', return_value=1), \
mock.patch.object(utils.LocalStorage, 'get_block_size', return_value=4096), \
mock.patch.object(utils.RemoteStorage, 'free_space'), \
mock.patch.object(utils.RemoteStorage, '__init__', return_value=None), \
mock.patch.object(utils.RemoteStorage, '__del__', return_value=None):
with self.assertRaises(celery.exceptions.Retry):
tasks.publish((1, [file_name])) # pylint: disable=no-value-for-parameter
mock_remove.assert_called()