'Create dagruns and esnure only ones with correct states are reset.'

def test_reset_orphaned_tasks_with_orphans(self):

    prefix = u'scheduler_job_test_test_reset_orphaned_tasks'
    states = [State.QUEUED, State.SCHEDULED, State.NONE, State.RUNNING, State.SUCCESS]
    states_to_reset = [State.QUEUED, State.SCHEDULED, State.NONE]
    dag = DAG(dag_id=prefix, start_date=DEFAULT_DATE, schedule_interval=u'@daily')
    tasks = []
    for i in range(len(states)):
        task_id = u'{}_task_{}'.format(prefix, i)
        task = DummyOperator(task_id=task_id, dag=dag)
        tasks.append(task)
    scheduler = SchedulerJob(**self.default_scheduler_args)
    session = settings.Session()
    dr1 = scheduler.create_dag_run(dag)
    dr2 = scheduler.create_dag_run(dag)
    dr1.state = State.RUNNING
    dr2.state = State.SUCCESS
    session.merge(dr1)
    session.merge(dr2)
    session.commit()
    dr1_tis = []
    dr2_tis = []
    for (i, (task, state)) in enumerate(zip(tasks, states)):
        ti1 = TI(task, dr1.execution_date)
        ti2 = TI(task, dr2.execution_date)
        ti1.refresh_from_db()
        ti2.refresh_from_db()
        ti1.state = state
        ti2.state = state
        dr1_tis.append(ti1)
        dr2_tis.append(ti2)
        session.merge(ti1)
        session.merge(ti2)
        session.commit()
    self.assertEqual(2, len(scheduler.reset_state_for_orphaned_tasks(session=session)))
    for ti in (dr1_tis + dr2_tis):
        ti.refresh_from_db()
    for (state, ti) in zip(states, dr1_tis):
        if (state in states_to_reset):
            self.assertIsNone(ti.state)
        else:
            self.assertEqual(state, ti.state)
    for (state, ti) in zip(states, dr2_tis):
        self.assertEqual(state, ti.state)
    for (state, ti) in zip(states, dr1_tis):
        ti.state = state
    session.commit()
    scheduler.reset_state_for_orphaned_tasks(filter_by_dag_run=dr1, session=session)
    for (state, ti) in zip(states, dr2_tis):
        self.assertEqual(state, ti.state)
    session.close()
