def test_image_should_be_pulled_when_not_exists(self, client_mock):
instance = mock.Mock(autospec=Instance,
**{
'start.return_value': 0,
'stop.return_value': 0,
})
client_mock.pull.return_value = '/tmp/busybox_latest.sif'
client_mock.instance.return_value = instance
client_mock.execute.return_value = {
'return_code': 0,
'message': 'message'
}
task = SingularityOperator(task_id='task-id',
image="docker://busybox",
command="echo hello",
pull_folder="/tmp",
force_pull=True)
task.execute({})
client_mock.instance.assert_called_once_with("/tmp/busybox_latest.sif",
options=[],
args=None,
start=False)
client_mock.pull.assert_called_once_with("docker://busybox",
stream=True,
pull_folder="/tmp")
client_mock.execute.assert_called_once_with(mock.ANY,
"echo hello",
return_result=True)
def test_command_is_required(self, command):
task = SingularityOperator(task_id='task-id',
image="docker://busybox",
command=command)
with six.assertRaisesRegex(self, AirflowException,
"You must define a command."):
task.execute({})
def test_working_dir(self, working_dir, expected_working_dir, client_mock):
instance = mock.Mock(autospec=Instance,
**{
'start.return_value': 0,
'stop.return_value': 0,
})
client_mock.pull.return_value = 'docker://busybox'
client_mock.instance.return_value = instance
client_mock.execute.return_value = {
'return_code': 0,
'message': 'message'
}
task = SingularityOperator(task_id='task-id',
image="docker://busybox",
command="echo hello",
force_pull=True,
working_dir=working_dir)
task.execute({})
client_mock.instance.assert_called_once_with(
"docker://busybox",
options=expected_working_dir,
args=None,
start=False)
def test_execute(self, client_mock):
instance = mock.Mock(autospec=Instance,
**{
'start.return_value': 0,
'stop.return_value': 0,
})
client_mock.instance.return_value = instance
client_mock.execute.return_value = {
'return_code': 0,
'message': 'message'
}
task = SingularityOperator(task_id='task-id',
image="docker://busybox",
command="echo hello")
task.execute({})
client_mock.instance.assert_called_once_with("docker://busybox",
options=[],
args=None,
start=False)
client_mock.execute.assert_called_once_with(mock.ANY,
"echo hello",
return_result=True)
execute_args, _ = client_mock.execute.call_args
self.assertIs(execute_args[0], instance)
instance.start.assert_called_once_with()
instance.stop.assert_called_once_with()
def test_command_is_required(self, command):
task = SingularityOperator(task_id='task-id', image="docker://busybox", command=command)
with pytest.raises(AirflowException, match="You must define a command."):
task.execute({})