def setUp(self, aws_hook_mock):
self.aws_hook_mock = aws_hook_mock
self.ecs_operator_args = {
'task_id':
'task',
'task_definition':
't',
'cluster':
'c',
'overrides': {},
'aws_conn_id':
None,
'region_name':
'eu-west-1',
'group':
'group',
'placement_constraints': [{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
'network_configuration': {
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
}
}
self.ecs = ECSOperator(**self.ecs_operator_args)
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1')
def test_execute_without_failures(self, launch_type, tags, aws_hook_mock,
check_mock, wait_mock):
client_mock = aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
ecs = ECSOperator(launch_type=launch_type,
tags=tags,
**self.ecs_operator_args)
ecs.execute(None)
aws_hook_mock.return_value.get_client_type.assert_called_once_with(
'ecs', region_name='eu-west-1')
extend_args = {}
if launch_type == 'FARGATE':
extend_args['platformVersion'] = 'LATEST'
if tags:
extend_args['tags'] = [{
'key': k,
'value': v
} for (k, v) in tags.items()]
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType=launch_type,
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
},
**extend_args)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(
ecs.arn,
'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55'
)
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(
task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1')
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1',
group='group',
placement_constraints=[{
'expression':
'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
network_configuration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
})
def test_execute_without_failures(self, launch_type, aws_hook_mock,
check_mock, wait_mock):
client_mock = aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
ecs = ECSOperator(launch_type=launch_type, **self.ecs_operator_args)
ecs.execute(None)
aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs',
region_name='eu-west-1')
extend_args = {}
if launch_type == 'FARGATE':
extend_args['platformVersion'] = 'LATEST'
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType=launch_type,
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
},
**extend_args
)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(ecs.arn,
'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55')
def ecs_operator(self, dag):
return ECSOperator(
task_id="ecs_task",
dag=dag,
cluster="cluster",
task_definition="definition",
overrides={
"containerOverrides": [{
"command": ["some", "command"],
"cpu": 50
}]
},
)
def ecs_operator_with_env(self, dag):
return ECSOperator(
task_id="ecs_task",
dag=dag,
cluster="cluster",
task_definition="definition",
overrides={
"containerOverrides": [{
"command": ["some", "command"],
"cpu":
50,
"environment": [{
"name": "USER_KEY",
"value": "USER_VALUE"
}],
}]
},
)
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(
task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1',
group='group',
placement_constraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
network_configuration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
}
)
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs_operator_args = {
'task_id': 'task',
'task_definition': 't',
'cluster': 'c',
'overrides': {},
'aws_conn_id': None,
'region_name': 'eu-west-1',
'group': 'group',
'placement_constraints': [{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
'network_configuration': {
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
}
}
self.ecs = ECSOperator(**self.ecs_operator_args)
class TestECSOperator(unittest.TestCase):
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def setUp(self, aws_hook_mock):
self.aws_hook_mock = aws_hook_mock
self.ecs_operator_args = {
'task_id':
'task',
'task_definition':
't',
'cluster':
'c',
'overrides': {},
'aws_conn_id':
None,
'region_name':
'eu-west-1',
'group':
'group',
'placement_constraints': [{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
'network_configuration': {
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
}
}
self.ecs = ECSOperator(**self.ecs_operator_args)
def test_init(self):
self.assertEqual(self.ecs.region_name, 'eu-west-1')
self.assertEqual(self.ecs.task_definition, 't')
self.assertEqual(self.ecs.aws_conn_id, None)
self.assertEqual(self.ecs.cluster, 'c')
self.assertEqual(self.ecs.overrides, {})
self.assertEqual(self.ecs.hook, self.aws_hook_mock.return_value)
self.aws_hook_mock.assert_called_once_with(aws_conn_id=None)
def test_template_fields_overrides(self):
self.assertEqual(self.ecs.template_fields, ('overrides', ))
@parameterized.expand([['EC2'], ['FARGATE']])
@mock.patch.object(ECSOperator, '_wait_for_task_ended')
@mock.patch.object(ECSOperator, '_check_success_task')
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def test_execute_without_failures(self, launch_type, aws_hook_mock,
check_mock, wait_mock):
client_mock = aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
ecs = ECSOperator(launch_type=launch_type, **self.ecs_operator_args)
ecs.execute(None)
aws_hook_mock.return_value.get_client_type.assert_called_once_with(
'ecs', region_name='eu-west-1')
extend_args = {}
if launch_type == 'FARGATE':
extend_args['platformVersion'] = 'LATEST'
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType=launch_type,
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
},
**extend_args)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(
ecs.arn,
'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55'
)
def test_execute_with_failures(self):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
resp_failures = deepcopy(RESPONSE_WITHOUT_FAILURES)
resp_failures['failures'].append('dummy error')
client_mock.run_task.return_value = resp_failures
with self.assertRaises(AirflowException):
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with(
'ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab'],
}
})
def test_wait_end_tasks(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
self.ecs._wait_for_task_ended()
client_mock.get_waiter.assert_called_once_with('tasks_stopped')
client_mock.get_waiter.return_value.wait.assert_called_once_with(
cluster='c', tasks=['arn'])
self.assertEqual(
sys.maxsize,
client_mock.get_waiter.return_value.config.max_attempts)
def test_check_success_tasks_raises(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'lastStatus': 'STOPPED',
'exitCode': 1
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is not in success state ", str(e.exception))
self.assertIn("'name': 'foo'", str(e.exception))
self.assertIn("'lastStatus': 'STOPPED'", str(e.exception))
self.assertIn("'exitCode': 1", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_tasks_raises_pending(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'PENDING'
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is still pending ", str(e.exception))
self.assertIn("'name': 'container-name'", str(e.exception))
self.assertIn("'lastStatus': 'PENDING'", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_tasks_raises_multiple(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'exitCode': 1
}, {
'name': 'bar',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_host_terminated_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'stoppedReason':
'Host EC2 (instance i-1234567890abcdef) terminated.',
"containers": [{
"containerArn":
"arn:aws:ecs:us-east-1:012345678910:container/e1ed7aac-d9b2-4315-8726-d2432bf11868", # noqa: E501
"lastStatus":
"RUNNING",
"name":
"wordpress",
"taskArn":
"arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55" # noqa: E501
}],
"desiredStatus":
"STOPPED",
"lastStatus":
"STOPPED",
"taskArn":
"arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55", # noqa: E501
"taskDefinitionArn":
"arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:11" # noqa: E501
}]
}
with self.assertRaises(AirflowException) as e:
self.ecs._check_success_task()
self.assertIn(
"The task was stopped because the host instance terminated:",
str(e.exception))
self.assertIn("Host EC2 (", str(e.exception))
self.assertIn(") terminated", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_task_not_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
'awsvpcConfiguration': {
'assignPublicIp': 'ENABLED',
'subnets': [awsNetworkSubnet]
}
},
'awslogs_group': '/ecs/' + awsTaskDefinition,
'awslogs_stream_prefix': 'ecs/' + awsContainerName,
'overrides': {
'containerOverrides': [
{
'name': awsContainerName,
'memoryReservation': 500,
},
],
},
}
# Create ECS operators
task_list = []
for index in range(len(yearsToAnalyze)):
ecs_operator_args = copy.deepcopy(ecs_operator_args_template)
ecs_operator_args['overrides']['containerOverrides'][0]['command'] = [
'-y', str(yearsToAnalyze[index])
]
ecs_operator = ECSOperator(task_id=str(yearsToAnalyze[index]),
dag=dag,
retries=AIRFLOW_ECS_OPERATOR_RETRIES,
retry_delay=timedelta(seconds=10),
**ecs_operator_args)
task_list.append(ecs_operator)
ecs_operator_args_template = {
'aws_conn_id': 'aws_default',
'region_name': awsRegionName,
'launch_type': 'FARGATE',
'cluster': awsCluster,
'task_definition': awsTaskDefinition,
'network_configuration': {
'awsvpcConfiguration': {
'assignPublicIp': 'ENABLED',
'subnets': [awsNetworkSubnet]
}
},
'awslogs_group': '/ecs/' + awsTaskDefinition,
'awslogs_stream_prefix': 'ecs/' + awsContainerName,
'overrides': {
'containerOverrides': [
{
'name': awsContainerName,
'memoryReservation': 500,
},
],
},
}
ecs_operator_args = copy.deepcopy(ecs_operator_args_template)
ecs_operator = ECSOperator(task_id='run_ml',
dag=dag,
retries=AIRFLOW_ECS_OPERATOR_RETRIES,
retry_delay=timedelta(seconds=10),
**ecs_operator_args)
network_configuration = {
"awsvpcConfiguration": {
"subnets": subnet_ids,
"securityGroups": ["sg-07381accdf2c5b195"],
"assignPublicIp": "ENABLED",
}
}
# Operator作成
t1 = ECSOperator(
task_id="run_collector",
dag=dag,
task_definition=task_definition,
cluster=cluster,
overrides=overrides,
region_name="us-west-1",
launch_type="FARGATE",
network_configuration=network_configuration,
)
# t2 = ECSOperator(
# task_id='sleep',
# bash_command='sleep 5',
# retries=3,
# dag=dag)
#
# templated_command = """
# {% for i in range(5) %}
# echo "{{ ds }}"
# echo "{{ macros.ds_add(ds, 7)}}"
# echo "{{ params.my_param }}"
class TestECSOperator(unittest.TestCase):
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(
task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1',
group='group',
placement_constraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
network_configuration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
}
)
def test_init(self):
self.assertEqual(self.ecs.region_name, 'eu-west-1')
self.assertEqual(self.ecs.task_definition, 't')
self.assertEqual(self.ecs.aws_conn_id, None)
self.assertEqual(self.ecs.cluster, 'c')
self.assertEqual(self.ecs.overrides, {})
self.assertEqual(self.ecs.hook, self.aws_hook_mock.return_value)
self.aws_hook_mock.assert_called_once_with(aws_conn_id=None)
def test_template_fields_overrides(self):
self.assertEqual(self.ecs.template_fields, ('overrides',))
@mock.patch.object(ECSOperator, '_wait_for_task_ended')
@mock.patch.object(ECSOperator, '_check_success_task')
def test_execute_without_failures(self, check_mock, wait_mock):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
platformVersion='LATEST',
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
}
)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(self.ecs.arn, 'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55')
def test_execute_with_failures(self):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
resp_failures = deepcopy(RESPONSE_WITHOUT_FAILURES)
resp_failures['failures'].append('dummy error')
client_mock.run_task.return_value = resp_failures
with self.assertRaises(AirflowException):
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
platformVersion='LATEST',
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc']
}
}
)
def test_wait_end_tasks(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
self.ecs._wait_for_task_ended()
client_mock.get_waiter.assert_called_once_with('tasks_stopped')
client_mock.get_waiter.return_value.wait.assert_called_once_with(cluster='c', tasks=['arn'])
self.assertEquals(sys.maxsize, client_mock.get_waiter.return_value.config.max_attempts)
def test_check_success_tasks_raises(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'lastStatus': 'STOPPED',
'exitCode': 1
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is not in success state ", str(e.exception))
self.assertIn("'name': 'foo'", str(e.exception))
self.assertIn("'lastStatus': 'STOPPED'", str(e.exception))
self.assertIn("'exitCode': 1", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_pending(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'PENDING'
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is still pending ", str(e.exception))
self.assertIn("'name': 'container-name'", str(e.exception))
self.assertIn("'lastStatus': 'PENDING'", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_multiple(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'exitCode': 1
}, {
'name': 'bar',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_task_not_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
":databricks: Databricks Cluster Created with ID: {{ task_instance.xcom_pull(task_ids='create_databricks_cluster') }}",
dag=dag)
create_cluster_notify.set_upstream(create_cluster)
train_model = ECSOperator(
task_id="train_model",
task_definition='trainmodelriskassessment',
cluster='TalendECS',
aws_conn_id='aws_default',
overrides={
'containerOverrides': [
{
'name':
"trainmodelriskassessment",
'command': [
"--context_param DATABRICKS_ENDPOINT=XXX",
"--context_param DATABRICKS_TOKEN=XXX",
"--context_param DATABRICKS_CLUSTER_ID={{ task_instance.xcom_pull(task_ids='create_databricks_cluster') }}"
]
},
]
},
region_name='us-east-1',
launch_type='EC2',
dag=dag)
# define list of lobs we want to run for
# loop through the lob's we want to use to build up our dag
s3_list_files = S3ListOperator(task_id="s3_list_files",
class TestECSOperator(unittest.TestCase):
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1')
def test_init(self):
self.assertEqual(self.ecs.region_name, 'eu-west-1')
self.assertEqual(self.ecs.task_definition, 't')
self.assertEqual(self.ecs.aws_conn_id, None)
self.assertEqual(self.ecs.cluster, 'c')
self.assertEqual(self.ecs.overrides, {})
self.assertEqual(self.ecs.hook, self.aws_hook_mock.return_value)
self.aws_hook_mock.assert_called_once_with(aws_conn_id=None)
def test_template_fields_overrides(self):
self.assertEqual(self.ecs.template_fields, ('overrides', ))
@mock.patch.object(ECSOperator, '_wait_for_task_ended')
@mock.patch.object(ECSOperator, '_check_success_task')
def test_execute_without_failures(self, check_mock, wait_mock):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with(
'ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t')
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(
self.ecs.arn,
'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55'
)
def test_execute_with_failures(self):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
resp_failures = deepcopy(RESPONSE_WITHOUT_FAILURES)
resp_failures['failures'].append('dummy error')
client_mock.run_task.return_value = resp_failures
with self.assertRaises(AirflowException):
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with(
'ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t')
def test_wait_end_tasks(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
self.ecs._wait_for_task_ended()
client_mock.get_waiter.assert_called_once_with('tasks_stopped')
client_mock.get_waiter.return_value.wait.assert_called_once_with(
cluster='c', tasks=['arn'])
self.assertEquals(
sys.maxsize,
client_mock.get_waiter.return_value.config.max_attempts)
def test_check_success_tasks_raises(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'lastStatus': 'STOPPED',
'exitCode': 1
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is not in success state ", str(e.exception))
self.assertIn("'name': 'foo'", str(e.exception))
self.assertIn("'lastStatus': 'STOPPED'", str(e.exception))
self.assertIn("'exitCode': 1", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_tasks_raises_pending(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'PENDING'
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is still pending ", str(e.exception))
self.assertIn("'name': 'container-name'", str(e.exception))
self.assertIn("'lastStatus': 'PENDING'", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_tasks_raises_multiple(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'exitCode': 1
}, {
'name': 'bar',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
def test_check_success_task_not_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c',
tasks=['arn'])
class TestECSOperator(unittest.TestCase):
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs_operator_args = {
'task_id': 'task',
'task_definition': 't',
'cluster': 'c',
'overrides': {},
'aws_conn_id': None,
'region_name': 'eu-west-1',
'group': 'group',
'placement_constraints': [{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}],
'network_configuration': {
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
}
}
self.ecs = ECSOperator(**self.ecs_operator_args)
def test_init(self):
self.assertEqual(self.ecs.region_name, 'eu-west-1')
self.assertEqual(self.ecs.task_definition, 't')
self.assertEqual(self.ecs.aws_conn_id, None)
self.assertEqual(self.ecs.cluster, 'c')
self.assertEqual(self.ecs.overrides, {})
self.assertEqual(self.ecs.hook, self.aws_hook_mock.return_value)
self.aws_hook_mock.assert_called_once_with(aws_conn_id=None)
def test_template_fields_overrides(self):
self.assertEqual(self.ecs.template_fields, ('overrides',))
@parameterized.expand([
['EC2'],
['FARGATE']
])
@mock.patch.object(ECSOperator, '_wait_for_task_ended')
@mock.patch.object(ECSOperator, '_check_success_task')
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def test_execute_without_failures(self, launch_type, aws_hook_mock,
check_mock, wait_mock):
client_mock = aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
ecs = ECSOperator(launch_type=launch_type, **self.ecs_operator_args)
ecs.execute(None)
aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs',
region_name='eu-west-1')
extend_args = {}
if launch_type == 'FARGATE':
extend_args['platformVersion'] = 'LATEST'
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType=launch_type,
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab']
}
},
**extend_args
)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(ecs.arn,
'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55')
def test_execute_with_failures(self):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
resp_failures = deepcopy(RESPONSE_WITHOUT_FAILURES)
resp_failures['failures'].append('dummy error')
client_mock.run_task.return_value = resp_failures
with self.assertRaises(AirflowException):
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs',
region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
launchType='EC2',
overrides={},
startedBy=mock.ANY, # Can by 'airflow' or 'Airflow'
taskDefinition='t',
group='group',
placementConstraints=[
{
'expression': 'attribute:ecs.instance-type =~ t2.*',
'type': 'memberOf'
}
],
networkConfiguration={
'awsvpcConfiguration': {
'securityGroups': ['sg-123abc'],
'subnets': ['subnet-123456ab'],
}
}
)
def test_wait_end_tasks(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
self.ecs._wait_for_task_ended()
client_mock.get_waiter.assert_called_once_with('tasks_stopped')
client_mock.get_waiter.return_value.wait.assert_called_once_with(
cluster='c', tasks=['arn'])
self.assertEqual(
sys.maxsize, client_mock.get_waiter.return_value.config.max_attempts)
def test_check_success_tasks_raises(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'lastStatus': 'STOPPED',
'exitCode': 1
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is not in success state ", str(e.exception))
self.assertIn("'name': 'foo'", str(e.exception))
self.assertIn("'lastStatus': 'STOPPED'", str(e.exception))
self.assertIn("'exitCode': 1", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(
cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_pending(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'PENDING'
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
# Ordering of str(dict) is not guaranteed.
self.assertIn("This task is still pending ", str(e.exception))
self.assertIn("'name': 'container-name'", str(e.exception))
self.assertIn("'lastStatus': 'PENDING'", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(
cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_multiple(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'exitCode': 1
}, {
'name': 'bar',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(
cluster='c', tasks=['arn'])
def test_host_terminated_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'stoppedReason': 'Host EC2 (instance i-1234567890abcdef) terminated.',
"containers": [
{
"containerArn": "arn:aws:ecs:us-east-1:012345678910:container/e1ed7aac-d9b2-4315-8726-d2432bf11868", # noqa: E501
"lastStatus": "RUNNING",
"name": "wordpress",
"taskArn": "arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55" # noqa: E501
}
],
"desiredStatus": "STOPPED",
"lastStatus": "STOPPED",
"taskArn": "arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55", # noqa: E501
"taskDefinitionArn": "arn:aws:ecs:us-east-1:012345678910:task-definition/hello_world:11" # noqa: E501
}]
}
with self.assertRaises(AirflowException) as e:
self.ecs._check_success_task()
self.assertIn(
"The task was stopped because the host instance terminated:",
str(e.exception))
self.assertIn("Host EC2 (", str(e.exception))
self.assertIn(") terminated", str(e.exception))
client_mock.describe_tasks.assert_called_once_with(
cluster='c', tasks=['arn'])
def test_check_success_task_not_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(
cluster='c', tasks=['arn'])
(TimeSensor(task_id='hold_on', target_time=time(hour=3), dag=dag) >>
ECSOperator(
task_id='run_ecs',
task_definition='airflow-etl-adwords',
cluster='ecs',
overrides={
'containerOverrides': [{
'name':
'app',
'environment': [{
'name': 'EXECUTION_DATE',
'value': '{{ ds }}'
}, {
'name': 'RDS_ENDPOINT',
'value': 'XXX'
}, {
'name': 'RDS_USER',
'value': 'XXX'
}, {
'name': 'RDS_PASSWORD',
'value': 'XXX'
}, {
'name': 'RDS_DATABASE',
'value': 'XXX'
}]
}]
},
aws_conn_id=
'', # This is required or else it will throw an AttributeError ('NoneType' object has no attribute 'upper')
region_name='ap-northeast-1',
dag=dag))
class TestECSOperator(unittest.TestCase):
@mock.patch('airflow.contrib.operators.ecs_operator.AwsHook')
def setUp(self, aws_hook_mock):
configuration.load_test_config()
self.aws_hook_mock = aws_hook_mock
self.ecs = ECSOperator(
task_id='task',
task_definition='t',
cluster='c',
overrides={},
aws_conn_id=None,
region_name='eu-west-1')
def test_init(self):
self.assertEqual(self.ecs.region_name, 'eu-west-1')
self.assertEqual(self.ecs.task_definition, 't')
self.assertEqual(self.ecs.aws_conn_id, None)
self.assertEqual(self.ecs.cluster, 'c')
self.assertEqual(self.ecs.overrides, {})
self.assertEqual(self.ecs.hook, self.aws_hook_mock.return_value)
self.aws_hook_mock.assert_called_once_with(aws_conn_id=None)
def test_template_fields_overrides(self):
self.assertEqual(self.ecs.template_fields, ('overrides',))
@mock.patch.object(ECSOperator, '_wait_for_task_ended')
@mock.patch.object(ECSOperator, '_check_success_task')
def test_execute_without_failures(self, check_mock, wait_mock):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
overrides={},
startedBy='Airflow',
taskDefinition='t'
)
wait_mock.assert_called_once_with()
check_mock.assert_called_once_with()
self.assertEqual(self.ecs.arn, 'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55')
def test_execute_with_failures(self):
client_mock = self.aws_hook_mock.return_value.get_client_type.return_value
resp_failures = deepcopy(RESPONSE_WITHOUT_FAILURES)
resp_failures['failures'].append('dummy error')
client_mock.run_task.return_value = resp_failures
with self.assertRaises(AirflowException):
self.ecs.execute(None)
self.aws_hook_mock.return_value.get_client_type.assert_called_once_with('ecs', region_name='eu-west-1')
client_mock.run_task.assert_called_once_with(
cluster='c',
overrides={},
startedBy='Airflow',
taskDefinition='t'
)
def test_wait_end_tasks(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
self.ecs._wait_for_task_ended()
client_mock.get_waiter.assert_called_once_with('tasks_stopped')
client_mock.get_waiter.return_value.wait.assert_called_once_with(cluster='c', tasks=['arn'])
self.assertEquals(sys.maxint, client_mock.get_waiter.return_value.config.max_attempts)
def test_check_success_tasks_raises(self):
client_mock = mock.Mock()
self.ecs.arn = 'arn'
self.ecs.client = client_mock
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'lastStatus': 'STOPPED',
'exitCode': 1
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
self.assertEquals(str(e.exception), "This task is not in success state {'containers': [{'lastStatus': 'STOPPED', 'name': 'foo', 'exitCode': 1}]}")
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_pending(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'PENDING'
}]
}]
}
with self.assertRaises(Exception) as e:
self.ecs._check_success_task()
self.assertEquals(str(e.exception), "This task is still pending {'containers': [{'lastStatus': 'PENDING', 'name': 'container-name'}]}")
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_tasks_raises_mutliple(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'foo',
'exitCode': 1
}, {
'name': 'bar',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
def test_check_success_task_not_raises(self):
client_mock = mock.Mock()
self.ecs.client = client_mock
self.ecs.arn = 'arn'
client_mock.describe_tasks.return_value = {
'tasks': [{
'containers': [{
'name': 'container-name',
'lastStatus': 'STOPPED',
'exitCode': 0
}]
}]
}
self.ecs._check_success_task()
client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
network_configuration = {
'awsvpcConfiguration': {
'subnets':
['subnet-id-1', 'subnet-id-2'
], # these are the actual, ugly ID values, not friendly names
'securityGroups': ['sg-id'],
'assignPublicIp': 'DISABLED'
}
}
# Can optionally provide override values (including commands) to the container
# used in the Task Definition. As found here:
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.run_task
overrides = {
'containerOverrides': [{
'name': 'nginx', # this is the container name that's launched
'command': ['custom', 'command']
}]
}
ecs_task = ECSOperator(
task_id='ecs_task', # this is Airflow's task name, don't get confused!
task_definition='task_def_name', # actual task def name in your cluster
cluster='cluster_name',
region_name='us-east-1', # required, for some reason
launch_type='FARGATE',
network_configuration=network_configuration,
overrides={}, # or `overrides` mapping from above
aws_conn_id=None, # let boto3 pull attached Role
dag=dag)
bash_command="echo 1",
dag=dag)
post_task = BashOperator(task_id="post_dbt", bash_command="echo 0", dag=dag)
dbt_top_events = ECSOperator(
task_id="dbt_top_events",
dag=dag,
aws_conn_id="aws_ecs",
cluster="MyCluster",
task_definition="dbt-cdk",
launch_type="FARGATE",
overrides={
"containerOverrides": [
{
"name": "dbt-cdk-container",
"command": ["dbt", "run"],
},
],
},
network_configuration={
"awsvpcConfiguration": {
"securityGroups": ["<SECURITY_GROUP_ID>"],
"subnets": ["<SUBNET_ID>", "<SUBNET_ID>"],
},
},
awslogs_group="/ecs/dbt-cdk",
awslogs_stream_prefix="ecs/dbt-cdk-container",
)
bash_task >> dbt_top_events >> post_task
'awsvpcConfiguration': {
'securityGroups': ['sg-hijk', 'sg-abcd'],
'subnets': ['subnet-lmn'],
'assignPublicIp': "ENABLED"
}
},
overrides={
'containerOverrides': [
{
'name': "my-container",
'command': ["python", "myCode.py",
str(param)]
}
]
},
region_name="us-east-1")
ecs_args = get_ecs_operator_args("{{ dag_run.conf['name'] }}")
dag = DAG( DAG_NAME,
schedule_interval=None,
default_args=default_args)
start_process = DummyOperator(task_id="start_process", dag=dag)
fargate_task = ECSOperator(task_id="fargate_task", **ecs_args, dag=dag)
start_process >> fargate_task
catchup=False,
max_active_runs=1,
)
# t1, t2 and t3 are examples of tasks created by instantiating operators
setup_task = BashOperator(task_id="setup", bash_command="date", dag=dag)
finish_task = BashOperator(task_id="tear_down", bash_command="date", dag=dag)
np.random.seed(593)
for i in range(N_MODELS):
t = ECSOperator(
# ECS-specific args
task_definition="generic_task:6",
cluster="tims-cluster",
# the work goes in here
overrides={
"containerOverrides": [{
"command": ["sleep",
str(np.random.poisson(10, size=None))]
}]
},
aws_conn_id="tims_aws_account",
launch_type="FARGATE",
# general operator args
task_id="train_model_%d" % i,
retries=0,
dag=dag,
)
t.set_upstream(setup_task)
t.set_downstream(finish_task)
