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_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_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)
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.providers.singularity.operators.singularity import SingularityOperator
with DAG(
'singularity_sample',
default_args={'retries': 1},
schedule_interval=timedelta(minutes=10),
start_date=datetime(2021, 1, 1),
catchup=False,
) as dag:
t1 = BashOperator(task_id='print_date', bash_command='date')
t2 = BashOperator(task_id='sleep', bash_command='sleep 5', retries=3)
t3 = SingularityOperator(
command='/bin/sleep 30',
image='docker://busybox:1.30.1',
task_id='singularity_op_tester',
)
t4 = BashOperator(task_id='print_hello',
bash_command='echo "hello world!!!"')
t1 >> [t2, t3]
t3 >> t4
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({})