Skip to content

Remove unnecessary DB clear in test #43731

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ephraimbuddy wants to merge 3 commits into from
Closed

Remove unnecessary DB clear in test #43731

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
581b34c
Remove unnecessary DB clear in test
ephraimbuddy Nov 6, 2024
d31a177
fixup! Remove unnecessary DB clear in test
ephraimbuddy Nov 6, 2024
cb97870
test
ephraimbuddy Nov 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 34 additions & 31 deletions tests/jobs/test_scheduler_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3965,7 +3965,6 @@ def test_create_dag_runs_assets(self, session, dag_maker):
- That the run created is on QUEUED State
- That dag_model has next_dagrun
"""
clear_db_dags()

asset1 = Asset(uri="ds1")
asset2 = Asset(uri="ds2")
Expand Down Expand Up @@ -5843,36 +5842,40 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce
callback_requests[0].simple_task_instance = None
assert expected_failure_callback_requests[0] == callback_requests[0]

def test_cleanup_stale_dags(self):
clear_db_dags()
dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False)
with create_session() as session:
dag = dagbag.get_dag("test_example_bash_operator")
dag.sync_to_db()
dm = DagModel.get_current("test_example_bash_operator")
# Make it "stale".
dm.last_parsed_time = timezone.utcnow() - timedelta(minutes=11)
session.merge(dm)

# This one should remain active.
dag = dagbag.get_dag("test_start_date_scheduling")
dag.sync_to_db()

session.flush()

scheduler_job = Job(executor=MockExecutor())
self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
self.job_runner.processor_agent = mock.MagicMock()

active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
assert active_dag_count == 2

self.job_runner._cleanup_stale_dags(session)

session.flush()

active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
assert active_dag_count == 1
# def test_cleanup_stale_dags(self):
# clear_db_dags()
# dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False)
# with create_session() as session:
# dag = dagbag.get_dag("test_example_bash_operator")
# dag.sync_to_db()
# dm = DagModel.get_current("test_example_bash_operator")
# # Make it "stale".
# dm.last_parsed_time = timezone.utcnow() - timedelta(minutes=11)
# session.merge(dm)
#
# # This one should remain active.
# dag = dagbag.get_dag("test_start_date_scheduling")
# dag.sync_to_db()
#
# session.flush()
#
# scheduler_job = Job(executor=MockExecutor())
# self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
# self.job_runner.processor_agent = mock.MagicMock()
#
# active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
# assert active_dag_count == 2
#
# self.job_runner._cleanup_stale_dags(session)
#
# session.flush()
#
# active_dag_count = session.query(func.count(DagModel.dag_id)).filter(DagModel.is_active).scalar()
# assert active_dag_count == 1
# # Adding back the cleared DB dags
# non_serialized_dagbag = DagBag(read_dags_from_db=False, include_examples=False)
# non_serialized_dagbag.sync_to_db()
# self.dagbag = DagBag(read_dags_from_db=True)

@mock.patch.object(settings, "USE_JOB_SCHEDULE", False)
def run_scheduler_until_dagrun_terminal(self):
Expand Down