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, []))