ashb commented on code in PR #44898:
URL: https://github.com/apache/airflow/pull/44898#discussion_r1883767599
##########
tests/models/test_dagbag.py:
##########
@@ -655,154 +607,6 @@ def f():
assert "tmp_file.py" in dagbag.import_errors
assert "DagBag import timeout for" in caplog.text
- @patch("airflow.models.dagbag.DagBag.collect_dags")
- @patch("airflow.models.serialized_dag.SerializedDagModel.write_dag")
- @patch("airflow.models.dag.DAG.bulk_write_to_db")
- def test_sync_to_db_is_retried(self, mock_bulk_write_to_db,
mock_s10n_write_dag, mock_collect_dags):
- """Test that dagbag.sync_to_db is retried on OperationalError"""
-
- dagbag = DagBag("/dev/null")
- mock_dag = mock.MagicMock()
- dagbag.dags["mock_dag"] = mock_dag
-
- op_error = OperationalError(statement=mock.ANY, params=mock.ANY,
orig=mock.ANY)
-
- # Mock error for the first 2 tries and a successful third try
- side_effect = [op_error, op_error, mock.ANY]
-
- mock_bulk_write_to_db.side_effect = side_effect
-
- mock_session = mock.MagicMock()
- dagbag.sync_to_db(session=mock_session)
-
- # Test that 3 attempts were made to run 'DAG.bulk_write_to_db'
successfully
- mock_bulk_write_to_db.assert_has_calls(
- [
- mock.call(mock.ANY, processor_subdir=None, session=mock.ANY),
- mock.call(mock.ANY, processor_subdir=None, session=mock.ANY),
- mock.call(mock.ANY, processor_subdir=None, session=mock.ANY),
- ]
- )
- # Assert that rollback is called twice (i.e. whenever OperationalError
occurs)
- mock_session.rollback.assert_has_calls([mock.call(), mock.call()])
- # Check that 'SerializedDagModel.write_dag' is also called
- # Only called once since the other two times the
'DAG.bulk_write_to_db' error'd
- # and the session was roll-backed before even reaching
'SerializedDagModel.write_dag'
- mock_s10n_write_dag.assert_has_calls(
- [
- mock.call(
- mock_dag, min_update_interval=mock.ANY,
processor_subdir=None, session=mock_session
- ),
- ]
- )
-
-
@patch("airflow.models.dagbag.settings.MIN_SERIALIZED_DAG_UPDATE_INTERVAL", 5)
- @patch("airflow.models.dagbag.DagBag._sync_perm_for_dag")
- def test_sync_to_db_syncs_dag_specific_perms_on_update(self,
mock_sync_perm_for_dag):
- """
- Test that dagbag.sync_to_db will sync DAG specific permissions when a
DAG is
- new or updated
- """
- db_clean_up()
- session = settings.Session()
- with time_machine.travel(tz.datetime(2020, 1, 5, 0, 0, 0), tick=False)
as frozen_time:
- dagbag = DagBag(
- dag_folder=os.path.join(TEST_DAGS_FOLDER,
"test_example_bash_operator.py"),
- include_examples=False,
- )
-
- def _sync_to_db():
- mock_sync_perm_for_dag.reset_mock()
- frozen_time.shift(20)
- dagbag.sync_to_db(session=session)
-
- dag = dagbag.dags["test_example_bash_operator"]
- dag.sync_to_db()
- _sync_to_db()
- mock_sync_perm_for_dag.assert_called_once_with(dag,
session=session)
-
- # DAG isn't updated
- _sync_to_db()
- mock_sync_perm_for_dag.assert_not_called()
-
- # DAG is updated
- dag.tags = ["new_tag"]
- _sync_to_db()
- session.commit()
- mock_sync_perm_for_dag.assert_called_once_with(dag,
session=session)
-
- @patch("airflow.www.security_appless.ApplessAirflowSecurityManager")
- def test_sync_perm_for_dag(self, mock_security_manager):
Review Comment:
Moved to test_collection
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
