class BasicTestCase(TestCase):
RESULT_ERROR = {ExecutorProtocol.RESULT: ExecutorProtocol.RESULT_ERROR}
RESULT_OK = Response(ResponseStatus.OK.value, "OK")
DOWNLOAD_STARTED = Response(
ResponseStatus.OK.value, ExecutorProtocol.DOWNLOAD_STARTED
)
DOWNLOAD_IN_PROGRESS = Response(
ResponseStatus.OK.value, ExecutorProtocol.DOWNLOAD_IN_PROGRESS
)
DOWNLOAD_FINISHED = Response(
ResponseStatus.OK.value, ExecutorProtocol.DOWNLOAD_FINISHED
)
ACCESS_LOG = Response(ResponseStatus.OK.value, 1)
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_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 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_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_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 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(rules, command, extra_fields=None, expect_reply=True):
"""Coroutine used to mock send_manager_command. It accepts the list of
expected calls and replies. When the actual call deviates from these
list the message with status ExecutorProtocol.RESULT_ERROR is sent to the
executor.
"""
rule_index = rules[0]
rules[0] += 1
expected, response = rules[rule_index]
if command != expected:
return Response(ResponseStatus.ERROR.value, "ERROR")
else:
if callable(response):
return response()
else:
return response
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 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_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, []))
