def test_get_result_from_dag_steps_unknown(get_latest_steps, get_action_info,
                                           log_warning_patch):
    helper = action_helper.ActionsHelper(action_id='irrelevant')
    result = helper.get_result_from_dag_steps()
    assert result == 'unknown'
    assert get_latest_steps.called
    assert get_action_info.called
    # Each critical step that had an unknown state should log a warning:
    assert log_warning_patch.call_count == 1
def test__get_latest_steps(get_all_steps, get_action_info):
    helper = action_helper.ActionsHelper(action_id='irrelevant')
    latest_steps_dict = helper._get_latest_steps()
    assert latest_steps_dict['task_A']['try_number'] == 3
    assert latest_steps_dict['task_B']['try_number'] == 1
    assert latest_steps_dict['task_C']['try_number'] == 2
    assert latest_steps_dict['task_D']['try_number'] == 1
    assert latest_steps_dict['task_E']['try_number'] == 1
    assert get_all_steps.called
    assert get_action_info.called
def test_get_result_from_dag_steps_running(get_latest_steps, get_action_info):
    helper = action_helper.ActionsHelper(action_id='irrelevant')
    result = helper.get_result_from_dag_steps()
    assert result == 'pending'
    assert get_latest_steps.called
    assert get_action_info.called
def test_get_result_from_dag_steps_success(get_latest_steps, get_action_info):
    helper = action_helper.ActionsHelper(action_id='irrelevant')
    result = helper.get_result_from_dag_steps()
    assert result == 'successful'
    assert get_latest_steps.called
    assert get_action_info.called
def test_get_step():
    # Set up actions helper
    action_id = '01CPV581B0CM8C9CA0CFRNVPPY'  # id in db
    actions = yaml.safe_load("""
---
  - id: 01CPV581B0CM8C9CA0CFRNVPPY
    name: update_software
    parameters: {}
    dag_id: update_software
    dag_execution_date: 2018-09-07T23:18:04
    user: admin
    datetime: 2018-09-07 23:18:04.38546+00
    context_marker: 10447c79-b02c-4dfd-a9e8-1362842f029d
...
""")
    action_helper.ActionsHelper._get_action_db = lambda \
            self, action_id: actions[0]

    tasks = yaml.safe_load("""
---
  - task_id: armada_get_status
    dag_id: update_software.armada_build
    execution_date: 2018-09-07 23:18:04
    start_date: 2018-09-07 23:18:55.950298
    end_date: 2018-09-07 23:18:58.159597
    duration: 2.209299
    state: success
    try_number: 1
    hostname: airflow-worker-0.airflow-worker-discovery.ucp.svc.cluster.local
    unixname: airflow
    job_id: 11
    pool:
    queue: default
    priority_weight: 3
    operator: ArmadaGetStatusOperator
    queued_dttm:
    pid: 249
    max_tries: 0
  - task_id: armada_get_status
    dag_id: update_software.armada_build
    execution_date: 2018-09-07 23:18:04
    start_date: 2018-09-07 23:18:55.950298
    end_date: 2018-09-07 23:18:58.159597
    duration: 2.209299
    state: success
    try_number: 2
    hostname: airflow-worker-1.airflow-worker-discovery.ucp.svc.cluster.local
    unixname: airflow
    job_id: 12
    pool:
    queue: default
    priority_weight: 3
    operator: ArmadaGetStatusOperator
    queued_dttm:
    pid: 249
    max_tries: 0
  - task_id: armada_post_apply
    dag_id: update_software.armada_build
    execution_date: 2018-09-07 23:18:04
    start_date: 2018-09-07 23:48:25.884615
    end_date: 2018-09-07 23:48:50.552757
    duration: 24.668142
    state: success
    try_number: 2
    hostname: airflow-worker-0.airflow-worker-discovery.ucp.svc.cluster.local
    unixname: airflow
    job_id: 13
    pool:
    queue: default
    priority_weight: 2
    operator: ArmadaPostApplyOperator
    queued_dttm:
    pid: 329
    max_tries: 3
  - task_id: armada_get_releases
    dag_id: update_software.armada_build
    execution_date: 2018-09-07 23:18:04
    start_date: 2018-09-07 23:48:59.024696
    end_date: 2018-09-07 23:49:01.471963
    duration: 2.447267
    state: success
    try_number: 1
    hostname: airflow-worker-0.airflow-worker-discovery.ucp.svc.cluster.local
    unixname: airflow
    job_id: 14
    pool:
    queue: default
    priority_weight: 1
    operator: ArmadaGetReleasesOperator
    queued_dttm:
    pid: 365
    max_tries: 0
  - task_id: armada_build
    dag_id: update_software
    execution_date: 2018-09-07 23:18:04
    start_date: 2018-09-07 23:18:47.447987
    end_date: 2018-09-07 23:49:02.397515
    duration: 1814.949528
    state: success
    try_number: 1
    hostname: airflow-worker-0.airflow-worker-discovery.ucp.svc.cluster.local
    unixname: airflow
    job_id: 9
    pool:
    queue: default
    priority_weight: 5
    operator: SubDagOperator
    queued_dttm: 2018-09-07 23:18:45.772501
    pid: 221
    max_tries: 0
...
""")
    action_helper.ActionsHelper._get_tasks_db = lambda \
            self, dag_id, execution_date: tasks
    actions_helper = action_helper.ActionsHelper(action_id=action_id)

    # Retrieve step
    step_id = 'armada_get_status'  # task_id in db

    # test backward compatibility with no additional param
    step = actions_helper.get_step(step_id)
    assert (step['hostname'].startswith('airflow-worker-0'))

    # test explicit None
    try_number = None
    step = actions_helper.get_step(step_id, try_number)
    assert (step['hostname'].startswith('airflow-worker-0'))

    # test try_number associated with 0 worker
    try_number = 1
    step = actions_helper.get_step(step_id, try_number)
    assert (step['hostname'].startswith('airflow-worker-0'))

    # test try_number associated with 1 worker
    try_number = 2
    step = actions_helper.get_step(step_id, try_number)
    assert (step['hostname'].startswith('airflow-worker-1'))