You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
140 lines
4.4 KiB
140 lines
4.4 KiB
# Copyright 2017 AT&T Intellectual Property. All other rights reserved. |
|
# |
|
# Licensed under the Apache License, Version 2.0 (the "License"); |
|
# you may not use this file except in compliance with the License. |
|
# You may obtain a copy of the License at |
|
# |
|
# http://www.apache.org/licenses/LICENSE-2.0 |
|
# |
|
# Unless required by applicable law or agreed to in writing, software |
|
# distributed under the License is distributed on an "AS IS" BASIS, |
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
# See the License for the specific language governing permissions and |
|
# limitations under the License. |
|
from unittest import mock |
|
|
|
import pytest |
|
from shipyard_airflow.plugins import concurrency_check_operator as operator |
|
from shipyard_airflow.plugins.concurrency_check_operator import ( |
|
ConcurrencyCheckOperator |
|
) |
|
from airflow.exceptions import AirflowException |
|
|
|
COLOR_SETS = [ |
|
set(['blue', 'green']), |
|
set(['blue', 'purple']), |
|
set(['red', 'purple']), |
|
set(['red', 'orange']), |
|
set(['yellow', 'green']), |
|
set(['yellow', 'orange']), |
|
] |
|
|
|
CONFLICT_SET = set(['cow', 'monkey', 'chicken']) |
|
|
|
|
|
def test_find_conflicting_dag_set(): |
|
""" |
|
Ensure that the right values are determined by find_conflicting_dag_set |
|
""" |
|
# Should not be found in the default set - no conflicts |
|
assert operator.DAG_RUN_SELECT_RUNNING_SQL is not None |
|
assert not operator.find_conflicting_dag_set("this_is_completely_cheese") |
|
|
|
# Check for contents vs the COLOR_SETS |
|
not_in_green_response_set = set(['purple', 'red', 'orange']) |
|
response_set = operator.find_conflicting_dag_set( |
|
dag_name='green', conflicting_dag_sets=COLOR_SETS) |
|
assert 'blue' in response_set |
|
assert 'yellow' in response_set |
|
assert not_in_green_response_set.isdisjoint(response_set) |
|
assert len(response_set) == 2 |
|
|
|
|
|
def get_executing_dags_stub_running_twice(): |
|
return [ |
|
('buffalo', 'now'), |
|
('buffalo', 'earlier'), |
|
('squirrel', 'ages ago'), |
|
] |
|
|
|
|
|
def get_executing_dags_stub(): |
|
return [ |
|
('buffalo', 'now'), |
|
('chicken', 'earlier'), |
|
('monkey', 'ages ago'), |
|
] |
|
|
|
|
|
def get_executing_dags_stub_no_conflicts(): |
|
return [ |
|
('buffalo', 'now'), |
|
('hedgehog', 'earlier'), |
|
('panda', 'ages ago'), |
|
] |
|
|
|
|
|
def test_find_conflicting_dag(): |
|
""" |
|
Ensure that: |
|
1) responds with a found conflict |
|
2) responds None if there is no conflict |
|
3) responds with the dag_id_to_check being searched for if it is running |
|
more than once. |
|
""" |
|
cco = ConcurrencyCheckOperator( |
|
conflicting_dag_set=CONFLICT_SET, |
|
task_id='bogus') |
|
|
|
# no conflicts |
|
cco.get_executing_dags = get_executing_dags_stub_no_conflicts |
|
assert cco.find_conflicting_dag('buffalo') is None |
|
|
|
# self is running twice |
|
cco.get_executing_dags = get_executing_dags_stub_running_twice |
|
assert cco.find_conflicting_dag('buffalo') != 'squirrel' |
|
assert cco.find_conflicting_dag('buffalo') == 'buffalo' |
|
|
|
# a conflict from the list |
|
cco.get_executing_dags = get_executing_dags_stub |
|
assert cco.find_conflicting_dag('buffalo') != 'monkey' |
|
assert cco.find_conflicting_dag('buffalo') == 'chicken' |
|
|
|
|
|
@mock.patch('shipyard_airflow.plugins.concurrency_check_operator.XcomPusher') |
|
def test_execute_exception(xcom_pusher): |
|
""" |
|
Run the whole execute function for testing |
|
""" |
|
cco = ConcurrencyCheckOperator( |
|
conflicting_dag_set=CONFLICT_SET, |
|
task_id='bogus') |
|
# dag_id of cow should cause monkey to conflict. |
|
cco.check_dag_id = 'cow' |
|
cco.get_executing_dags = get_executing_dags_stub |
|
try: |
|
context = {'task_instance': None} |
|
cco.execute(context) |
|
pytest.fail('AirflowException should have been raised') |
|
except AirflowException as airflow_exception: |
|
assert 'Aborting run' in airflow_exception.args[0] |
|
|
|
|
|
@mock.patch('shipyard_airflow.plugins.concurrency_check_operator.XcomPusher') |
|
def test_execute_success(xcom_pusher): |
|
""" |
|
Run the whole execute function for testing - successfully! |
|
""" |
|
cco = ConcurrencyCheckOperator( |
|
conflicting_dag_set=set(['car', 'truck']), |
|
task_id='bogus') |
|
|
|
# dag_id of airplane should have no conflicts |
|
cco.check_dag_id = 'airplane' |
|
cco.get_executing_dags = get_executing_dags_stub |
|
try: |
|
context = {'task_instance': None} |
|
cco.execute(context) |
|
assert True |
|
except AirflowException: |
|
pytest.fail('AirflowException should not have been raised')
|
|
|