Lee-W commented on code in PR #69977:
URL: https://github.com/apache/airflow/pull/69977#discussion_r3600191143
##########
airflow-core/tests/unit/assets/test_manager.py:
##########
@@ -211,6 +211,48 @@ def test_register_asset_change_no_downstreams(self,
session, mock_task_instance)
)
assert
session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 0
+ @pytest.mark.parametrize(
+ ("dialect_name", "expected_helper"),
+ [
+ ("postgresql", "_queue_dagruns_nonpartitioned_postgres"),
+ ("mysql", "_queue_dagruns_nonpartitioned_mysql"),
+ ("sqlite", "_queue_dagruns_nonpartitioned_slow_path"),
+ ],
+ )
+ def test_queue_dagruns_routes_by_dialect(self, dialect_name,
expected_helper):
+ """Test that _queue_dagruns routes to the dialect-appropriate queue
helper."""
+ dag = DagModel(dag_id="dag1")
+ session = mock.MagicMock(spec=Session)
+ with (
+ mock.patch("airflow.assets.manager.get_dialect_name",
return_value=dialect_name),
+ mock.patch.object(AssetManager, "_queue_partitioned_dags"),
+ mock.patch.object(AssetManager, expected_helper) as mock_helper,
+ ):
+ AssetManager._queue_dagruns(
+ asset_id=1,
+ dags_to_queue={dag},
+ partition_key=None,
+ partition_date=None,
+ event=mock.MagicMock(),
+ task_instance=None,
+ session=session,
+ )
+ mock_helper.assert_called_once_with(1, {dag}, session)
+
+ def test_queue_dagruns_nonpartitioned_mysql_builds_upsert(self):
+ """Test that the MySQL queue path emits an INSERT ... ON DUPLICATE KEY
UPDATE."""
+ from sqlalchemy.dialects import mysql
Review Comment:
probably don't need to make it an inline import?
--
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]
