ashb commented on code in PR #44898:
URL: https://github.com/apache/airflow/pull/44898#discussion_r1883768289
##########
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):
- """
- Test that dagbag._sync_perm_for_dag will call
ApplessAirflowSecurityManager.sync_perm_for_dag
- """
- db_clean_up()
- with create_session() as session:
- security_manager = ApplessAirflowSecurityManager(session)
- mock_sync_perm_for_dag =
mock_security_manager.return_value.sync_perm_for_dag
- mock_sync_perm_for_dag.side_effect =
security_manager.sync_perm_for_dag
-
- dagbag = DagBag(
- dag_folder=os.path.join(TEST_DAGS_FOLDER,
"test_example_bash_operator.py"),
- include_examples=False,
- )
- dag = dagbag.dags["test_example_bash_operator"]
-
- def _sync_perms():
- mock_sync_perm_for_dag.reset_mock()
- DagBag._sync_perm_for_dag(dag, session=session)
-
- # perms dont exist
- _sync_perms()
-
mock_sync_perm_for_dag.assert_called_once_with("test_example_bash_operator",
None)
-
- # perms now exist
- _sync_perms()
-
mock_sync_perm_for_dag.assert_called_once_with("test_example_bash_operator",
None)
-
- # Always sync if we have access_control
- dag.access_control = {"Public": {"can_read"}}
- _sync_perms()
- mock_sync_perm_for_dag.assert_called_once_with(
- "test_example_bash_operator", {"Public": {"DAGs":
{"can_read"}}}
- )
-
- @patch("airflow.www.security_appless.ApplessAirflowSecurityManager")
- def test_sync_perm_for_dag_with_dict_access_control(self,
mock_security_manager):
Review Comment:
This was not moved --it was basically the same as test_sync_perm_for_dag but
with a different format for the perms -- that is handled inside DAG though, so
not relevant to test again here.
--
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]
