'Test that task retries are handled properly'

@patch.object(TI, u'pool_full')
def test_retry_handling(self, mock_pool_full):

    mock_pool_full.return_value = False
    dag = models.DAG(dag_id=u'test_retry_handling')
    task = BashOperator(task_id=u'test_retry_handling_op', bash_command=u'exit    1', retries=1, retry_delay=datetime.timedelta(seconds=0), dag=dag, owner=u'airflow', start_date=datetime.datetime(2016, 2, 1, 0, 0, 0))
    def run_with_error(ti):
        try:
            ti.run()
        except AirflowException:
            pass
    ti = TI(task=task, execution_date=datetime.datetime.now())
    run_with_error(ti)
    self.assertEqual(ti.state, State.UP_FOR_RETRY)
    self.assertEqual(ti.try_number, 1)
    run_with_error(ti)
    self.assertEqual(ti.state, State.FAILED)
    self.assertEqual(ti.try_number, 2)
    dag.clear()
    run_with_error(ti)
    self.assertEqual(ti.state, State.UP_FOR_RETRY)
    self.assertEqual(ti.try_number, 3)
    run_with_error(ti)
    self.assertEqual(ti.state, State.FAILED)
    self.assertEqual(ti.try_number, 4)
