def setUp(self):
self.communicator_mock = MagicMock()
self.missing_data = [
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
"to_connector": "local",
}
]
self.DOWNLOAD_STARTED_LOCK = Message.command(
"download_started",
{
ExecutorProtocol.STORAGE_LOCATION_ID: 2,
ExecutorProtocol.DOWNLOAD_STARTED_LOCK: True,
},
)
self.DOWNLOAD_STARTED_NO_LOCK = Message.command(
"download_started",
{
ExecutorProtocol.STORAGE_LOCATION_ID: 2,
ExecutorProtocol.DOWNLOAD_STARTED_LOCK: False,
},
)
self.MISSING_DATA = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
self.MISSING_DATA_RESPONSE = Response(
ResponseStatus.OK.value, self.missing_data.copy()
)
return super().setUp()
def test_handle_get_referenced_files(self):
obj = Message.command(ExecutorProtocol.GET_REFERENCED_FILES, "")
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name="local",
status=StorageLocation.STATUS_DONE,
url=str(self.file_storage.id),
)
path = Path(storage_location.get_path(filename="output.txt"))
path.parent.mkdir(exist_ok=True, parents=True)
path.touch()
data = Data.objects.get(id=1)
data.process.output_schema = [{
"name": "output_file",
"type": "basic:file:"
}]
data.process.save()
data.output = {"output_file": {"file": "output.txt"}}
data.save()
response = self.processor.handle_get_referenced_files(
obj, self.manager)
expected = Response(
ResponseStatus.OK.value,
[
"jsonout.txt",
"stdout.txt",
"output.txt",
],
)
self.assertEqual(response, expected)
def test_handle_missing_data_locations(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(parent=parent,
child=child,
kind=DataDependency.KIND_IO)
storage_location = StorageLocation.objects.create(
file_storage=parent.location,
connector_name="not_local",
status=StorageLocation.STATUS_DONE,
url="url",
)
response = self.processor.handle_missing_data_locations(
obj, self.manager)
self.assertEqual(StorageLocation.all_objects.count(), 3)
created = StorageLocation.all_objects.last()
expected = Response(
ResponseStatus.OK.value,
{
"url": {
"data_id": parent.id,
"from_connector": "not_local",
"from_storage_location_id": storage_location.id,
"to_storage_location_id": created.id,
"to_connector": "local",
}
},
)
self.assertEqual(response, expected)
def test_transfer_downloadmulti(self):
self.missing_data = [
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
},
{
"connector_name": "CONNECTOR",
"url": "URL",
"data_id": "1",
"from_storage_location_id": 1,
"to_storage_location_id": 2,
},
]
self.MISSING_DATA_RESPONSE = Response(
ResponseStatus.OK.value, self.missing_data.copy()
)
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_IN_PROGRESS),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
(self.DOWNLOAD_STARTED_NO_LOCK, self.DOWNLOAD_FINISHED),
]
download_command = MagicMock(side_effect=lambda a, b: coroutine(True))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertTrue(result)
download_command.assert_called_once()
def test_handle_download_started_no_location(self):
obj = Message.command(
ExecutorProtocol.DOWNLOAD_STARTED,
{
"storage_location_id": -2,
"download_started_lock": True,
},
)
with self.assertRaises(StorageLocation.DoesNotExist):
self.processor.handle_download_started(obj, self.manager)
def test_handle_download_aborted(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name="local",
status=StorageLocation.STATUS_UPLOADING,
)
obj = Message.command(ExecutorProtocol.DOWNLOAD_ABORTED, storage_location.id)
self.processor.handle_download_aborted(obj)
storage_location.refresh_from_db()
self.assertEqual(storage_location.status, StorageLocation.STATUS_PREPARING)
def test_handle_missing_data_locations_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(
parent=parent, child=child, kind=DataDependency.KIND_IO
)
response = self.processor.handle_missing_data_locations(obj)
expected = Response(ResponseStatus.ERROR.value, "No storage location exists")
self.assertEqual(response, expected)
self.assertEqual(StorageLocation.all_objects.count(), 1)
def setUp(self):
self.communicator_mock = MagicMock()
self.missing_data = {
"connector_name": "S3",
"url": "transfer_url",
"data_id": 1,
"from_storage_location_id": 1,
"to_storage_location_id": 2,
}
self.COMMAND_DOWNLOAD_FINISHED = Message.command(
ExecutorProtocol.DOWNLOAD_FINISHED, 2
)
self.COMMAND_DOWNLOAD_ABORTED = Message.command(
ExecutorProtocol.DOWNLOAD_ABORTED, 2
)
self.COMMAND_GET_FILES = Message.command(
ExecutorProtocol.GET_FILES_TO_DOWNLOAD, 1
)
self.FILES_LIST = Response(ResponseStatus.OK.value, ["1", "dir/1"])
return super().setUp()
def test_transfer_downloaded(self):
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_FINISHED),
]
download_command = MagicMock(return_value=coroutine(True))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertTrue(result)
download_command.assert_not_called()
def test_no_transfer(self):
self.MISSING_DATA_RESPONSE = Response(ResponseStatus.OK.value, [])
send_command = MagicMock(
side_effect=[
coroutine(self.MISSING_DATA_RESPONSE),
coroutine(self.RESULT_OK),
]
)
self.communicator_mock.send_command = send_command
run_async(transfer._transfer_data(self.communicator_mock))
send_command.assert_called_once_with(
Message.command("missing_data_locations", "")
)
async def send_single_message():
"""Open connection to listener and send single message."""
connection_string = f"{protocol}://{host}:{port}"
zmq_context = zmq.asyncio.Context.instance()
zmq_socket = zmq_context.socket(zmq.DEALER)
zmq_socket.setsockopt(zmq.IDENTITY, b"1")
zmq_socket.connect(connection_string)
communicator = ZMQCommunicator(zmq_socket,
"init_container <-> listener",
logger)
async with communicator:
await asyncio.ensure_future(
communicator.send_command(
Message.command("update_status", "PP")))
def test_handle_resolve_data_path(self):
"""Test data path resolwing."""
data = Data.objects.get(id=1)
Worker.objects.get_or_create(data=data, status=Worker.STATUS_PREPARING)
message = Message.command("resolve_data_path", data.pk)
response = self.manager.process_command(message)
assert response.message_data == str(constants.INPUTS_VOLUME)
connector_name = "local"
self.storage_location = StorageLocation.objects.create(
file_storage=self.file_storage,
connector_name=connector_name,
status="OK")
response = self.manager.process_command(message)
self.assertEqual(response.message_data, f"/data_{connector_name}")
def test_transfer_failed(self):
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
]
download_command = MagicMock(return_value=coroutine(False))
send_command = MagicMock(side_effect=partial(send, commands))
result = self._test_workflow(send_command, download_command, commands, True)
self.assertFalse(result)
download_command.assert_called_once_with(
self.missing_data[0], self.communicator_mock
)
async def terminate(self):
"""Send the terminate command to the worker.
Peer should terminate by itself and send finish message back to us.
"""
await database_sync_to_async(self._save_error
)("Processing was cancelled.")
# Ignore the possible timeout.
with suppress(RuntimeError):
await self._listener.communicator.send_command(
Message.command("terminate", ""),
peer_identity=str(self.data_id).encode(),
response_timeout=1,
)
def test_transfer_protocol_fail(self):
async def raise_exception(*args, **kwargs):
raise RuntimeError("Protocol error")
commands = [
1,
(self.MISSING_DATA, self.MISSING_DATA_RESPONSE),
(Message.command("update_status", "PP"), self.RESULT_OK),
(self.DOWNLOAD_STARTED_LOCK, self.DOWNLOAD_STARTED),
]
download_command = MagicMock(return_value=coroutine(True))
send_command = MagicMock(side_effect=raise_exception)
with self.assertRaises(RuntimeError):
self._test_workflow(send_command, download_command, commands, True)
download_command.assert_not_called()
def test_handle_download_started_ok_no_lock_preparing(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage, connector_name="local")
obj = Message.command(
ExecutorProtocol.DOWNLOAD_STARTED,
{
"storage_location_id": storage_location.id,
"download_started_lock": False,
},
)
response = self.processor.handle_download_started(obj, self.manager)
self.assertEqual(response,
Response(ResponseStatus.OK.value, "download_started"))
storage_location.refresh_from_db()
self.assertEqual(storage_location.status,
StorageLocation.STATUS_PREPARING)
def test_handle_missing_data_locations_none(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
parent = Data.objects.get(id=2)
child = Data.objects.get(id=1)
DataDependency.objects.create(
parent=parent, child=child, kind=DataDependency.KIND_IO
)
StorageLocation.objects.create(
file_storage=parent.location,
connector_name="local",
status=StorageLocation.STATUS_DONE,
url="url",
)
response = self.processor.handle_missing_data_locations(obj)
expected = Response(ResponseStatus.OK.value, [])
self.assertEqual(response, expected)
self.assertEqual(StorageLocation.all_objects.count(), 2)
def test_handle_get_files_to_download(self):
obj = Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD,
self.storage_location.id)
response = self.processor.handle_get_files_to_download(
obj, self.manager)
expected = Response(
ResponseStatus.OK.value,
[{
"id": self.path.id,
"path": "test.me",
"size": -1,
"md5": "md5",
"crc32c": "crc",
"awss3etag": "aws",
"chunk_size": BaseStorageConnector.CHUNK_SIZE,
}],
)
self.assertEqual(response, expected)
def test_handle_download_finished(self):
storage_location = StorageLocation.objects.create(
file_storage=self.file_storage, connector_name="local"
)
obj = Message.command(ExecutorProtocol.DOWNLOAD_FINISHED, storage_location.id)
with patch(
"resolwe.storage.models.FileStorage.default_storage_location",
self.storage_location,
):
response = self.processor.handle_download_finished(obj)
self.assertEqual(response.response_status, ResponseStatus.OK)
storage_location.refresh_from_db()
self.assertEqual(storage_location.status, StorageLocation.STATUS_DONE)
self.assertEqual(storage_location.files.count(), 1)
file = storage_location.files.get()
self.assertEqual(file.path, "test.me")
self.assertEqual(file.md5, "md5")
self.assertEqual(file.crc32c, "crc")
self.assertEqual(file.awss3etag, "aws")
def test_handle_download_aborted_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.DOWNLOAD_ABORTED, -2)
response = self.processor.handle_download_aborted(obj, self.manager)
self.assertEqual(response.response_status, ResponseStatus.OK)
def test_handle_download_finished_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.DOWNLOAD_FINISHED, -2)
with self.assertRaises(StorageLocation.DoesNotExist):
self.processor.handle_download_finished(obj, self.manager)
def test_handle_missing_data_locations_missing_data(self):
obj = Message.command(ExecutorProtocol.MISSING_DATA_LOCATIONS, "")
response = self.processor.handle_missing_data_locations(
obj, self.manager)
self.assertEqual(response, Response(ResponseStatus.OK.value, {}))
def test_handle_get_files_to_download_missing_storage_location(self):
obj = Message.command(ExecutorProtocol.GET_FILES_TO_DOWNLOAD, -2)
response = self.processor.handle_get_files_to_download(
obj, self.manager)
self.assertEqual(response, Response(ResponseStatus.OK.value, []))