def test_dispatch_execute_ignore(mock_write_to_file, mock_upload_dir,
mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
ctx = context_manager.FlyteContext.current_context()
# IgnoreOutputs
with ctx.new_execution_context(
mode=context_manager.ExecutionState.Mode.TASK_EXECUTION) as ctx:
python_task = mock.MagicMock()
python_task.dispatch_execute.side_effect = IgnoreOutputs()
empty_literal_map = _literal_models.LiteralMap({}).to_flyte_idl()
mock_load_proto.return_value = empty_literal_map
_dispatch_execute(ctx, python_task, "inputs path", "outputs prefix")
assert mock_write_to_file.call_count == 0
def test_dispatch_execute_exception(mock_write_to_file, mock_upload_dir,
mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
ctx = context_manager.FlyteContext.current_context()
with ctx.new_execution_context(
mode=context_manager.ExecutionState.Mode.TASK_EXECUTION) as ctx:
python_task = mock.MagicMock()
python_task.dispatch_execute.side_effect = Exception("random")
empty_literal_map = _literal_models.LiteralMap({}).to_flyte_idl()
mock_load_proto.return_value = empty_literal_map
def verify_output(*args, **kwargs):
assert isinstance(args[0], ErrorDocument)
mock_write_to_file.side_effect = verify_output
_dispatch_execute(ctx, python_task, "inputs path", "outputs prefix")
assert mock_write_to_file.call_count == 1
def test_dispatch_execute_void(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
python_task = mock.MagicMock()
python_task.dispatch_execute.return_value = VoidPromise("testing")
empty_literal_map = _literal_models.LiteralMap({}).to_flyte_idl()
mock_load_proto.return_value = empty_literal_map
def verify_output(*args, **kwargs):
assert args[0] == empty_literal_map
mock_write_to_file.side_effect = verify_output
_dispatch_execute(ctx, python_task, "inputs path", "outputs prefix")
assert mock_write_to_file.call_count == 1