Lee-W commented on code in PR #62184:
URL: https://github.com/apache/airflow/pull/62184#discussion_r3014847626
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -196,6 +211,131 @@ def test_custom_policy(self):
settings.task_must_have_owners(task_instance)
+class TestMetadataEngineHooks:
+ """Tests for the overridable create_metadata_engine /
create_async_metadata_engine hooks."""
+
+ def setup_method(self):
+ self.old_modules = dict(sys.modules)
+ from airflow import settings
+
+ self._orig_create_metadata_engine = settings.create_metadata_engine
+ self._orig_create_async_metadata_engine =
settings.create_async_metadata_engine
+
+ def teardown_method(self):
+ from airflow import settings
+
+ settings.create_metadata_engine = self._orig_create_metadata_engine
+ settings.create_async_metadata_engine =
self._orig_create_async_metadata_engine
+ for mod in [m for m in sys.modules if m not in self.old_modules]:
+ del sys.modules[mod]
+
+ @patch("airflow.settings.create_metadata_engine")
+ @patch("airflow.settings._configure_async_session")
+ def test_configure_orm_delegates_to_create_metadata_engine(self,
mock_async_session, mock_create_engine):
+ """configure_orm() must call create_metadata_engine, not create_engine
directly."""
+ from airflow import settings
+
+ mock_create_engine.return_value = MagicMock()
+
+ with (
+ patch("os.environ", {"_AIRFLOW_SKIP_DB_TESTS": "false"}),
+ patch("airflow.settings.SQL_ALCHEMY_CONN", "sqlite://"),
+ patch("airflow.settings.Session"),
+ patch("airflow.settings.engine"),
+ patch("airflow.settings.setup_event_handlers"),
+ patch("airflow.settings.mask_secret", create=True),
+ patch("airflow._shared.secrets_masker.mask_secret"),
+ ):
+ settings.configure_orm()
+
+ mock_create_engine.assert_called_once()
+ mock_async_session.assert_called_once()
Review Comment:
```suggestion
assert mock_create_engine.mock_calls = [call(...)]
assert mock_async_session.mock_calls = [call(...)]
```
would be better if we can check the args this way
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -196,6 +211,131 @@ def test_custom_policy(self):
settings.task_must_have_owners(task_instance)
+class TestMetadataEngineHooks:
+ """Tests for the overridable create_metadata_engine /
create_async_metadata_engine hooks."""
+
+ def setup_method(self):
+ self.old_modules = dict(sys.modules)
+ from airflow import settings
+
+ self._orig_create_metadata_engine = settings.create_metadata_engine
+ self._orig_create_async_metadata_engine =
settings.create_async_metadata_engine
+
+ def teardown_method(self):
+ from airflow import settings
+
+ settings.create_metadata_engine = self._orig_create_metadata_engine
+ settings.create_async_metadata_engine =
self._orig_create_async_metadata_engine
+ for mod in [m for m in sys.modules if m not in self.old_modules]:
+ del sys.modules[mod]
+
+ @patch("airflow.settings.create_metadata_engine")
+ @patch("airflow.settings._configure_async_session")
+ def test_configure_orm_delegates_to_create_metadata_engine(self,
mock_async_session, mock_create_engine):
+ """configure_orm() must call create_metadata_engine, not create_engine
directly."""
+ from airflow import settings
+
+ mock_create_engine.return_value = MagicMock()
+
+ with (
+ patch("os.environ", {"_AIRFLOW_SKIP_DB_TESTS": "false"}),
+ patch("airflow.settings.SQL_ALCHEMY_CONN", "sqlite://"),
+ patch("airflow.settings.Session"),
+ patch("airflow.settings.engine"),
+ patch("airflow.settings.setup_event_handlers"),
+ patch("airflow.settings.mask_secret", create=True),
+ patch("airflow._shared.secrets_masker.mask_secret"),
+ ):
+ settings.configure_orm()
+
+ mock_create_engine.assert_called_once()
+ mock_async_session.assert_called_once()
+ call_kwargs = mock_create_engine.call_args
+ assert call_kwargs[0][0] == "sqlite://"
+ assert "engine_args" in call_kwargs[1]
+ assert "connect_args" in call_kwargs[1]
+
+ @patch("airflow.settings.create_async_metadata_engine")
+ def test_configure_async_session_delegates_to_create_async_metadata_engine(
+ self, mock_create_async_engine
+ ):
+ """_configure_async_session() must call
create_async_metadata_engine."""
+ from airflow import settings
+
+ mock_create_async_engine.return_value = MagicMock()
+
+ with patch("airflow.settings.SQL_ALCHEMY_CONN_ASYNC",
"sqlite+aiosqlite://"):
+ settings._configure_async_session()
+
+ mock_create_async_engine.assert_called_once()
+ call_kwargs = mock_create_async_engine.call_args
+ assert call_kwargs[0][0] == "sqlite+aiosqlite://"
+ assert "connect_args" in call_kwargs[1]
+
+ @patch("airflow.settings.create_async_metadata_engine")
+ def test_configure_async_session_skips_when_no_async_conn(self,
mock_create_async_engine):
+ """_configure_async_session() must not call the hook when
SQL_ALCHEMY_CONN_ASYNC is empty."""
+ from airflow import settings
+
+ with patch("airflow.settings.SQL_ALCHEMY_CONN_ASYNC", ""):
+ settings._configure_async_session()
+
+ mock_create_async_engine.assert_not_called()
Review Comment:
```suggestion
assert mock_async_session.mock_calls == []
```
##########
airflow-core/tests/unit/core/test_settings.py:
##########
@@ -196,6 +211,131 @@ def test_custom_policy(self):
settings.task_must_have_owners(task_instance)
+class TestMetadataEngineHooks:
+ """Tests for the overridable create_metadata_engine /
create_async_metadata_engine hooks."""
+
+ def setup_method(self):
+ self.old_modules = dict(sys.modules)
+ from airflow import settings
+
+ self._orig_create_metadata_engine = settings.create_metadata_engine
+ self._orig_create_async_metadata_engine =
settings.create_async_metadata_engine
+
+ def teardown_method(self):
+ from airflow import settings
+
+ settings.create_metadata_engine = self._orig_create_metadata_engine
+ settings.create_async_metadata_engine =
self._orig_create_async_metadata_engine
+ for mod in [m for m in sys.modules if m not in self.old_modules]:
+ del sys.modules[mod]
+
+ @patch("airflow.settings.create_metadata_engine")
+ @patch("airflow.settings._configure_async_session")
+ def test_configure_orm_delegates_to_create_metadata_engine(self,
mock_async_session, mock_create_engine):
+ """configure_orm() must call create_metadata_engine, not create_engine
directly."""
+ from airflow import settings
+
+ mock_create_engine.return_value = MagicMock()
+
+ with (
+ patch("os.environ", {"_AIRFLOW_SKIP_DB_TESTS": "false"}),
+ patch("airflow.settings.SQL_ALCHEMY_CONN", "sqlite://"),
+ patch("airflow.settings.Session"),
+ patch("airflow.settings.engine"),
+ patch("airflow.settings.setup_event_handlers"),
+ patch("airflow.settings.mask_secret", create=True),
+ patch("airflow._shared.secrets_masker.mask_secret"),
+ ):
+ settings.configure_orm()
+
+ mock_create_engine.assert_called_once()
+ mock_async_session.assert_called_once()
+ call_kwargs = mock_create_engine.call_args
+ assert call_kwargs[0][0] == "sqlite://"
+ assert "engine_args" in call_kwargs[1]
+ assert "connect_args" in call_kwargs[1]
+
+ @patch("airflow.settings.create_async_metadata_engine")
+ def test_configure_async_session_delegates_to_create_async_metadata_engine(
+ self, mock_create_async_engine
+ ):
+ """_configure_async_session() must call
create_async_metadata_engine."""
+ from airflow import settings
+
+ mock_create_async_engine.return_value = MagicMock()
+
+ with patch("airflow.settings.SQL_ALCHEMY_CONN_ASYNC",
"sqlite+aiosqlite://"):
+ settings._configure_async_session()
+
+ mock_create_async_engine.assert_called_once()
+ call_kwargs = mock_create_async_engine.call_args
+ assert call_kwargs[0][0] == "sqlite+aiosqlite://"
+ assert "connect_args" in call_kwargs[1]
+
+ @patch("airflow.settings.create_async_metadata_engine")
+ def test_configure_async_session_skips_when_no_async_conn(self,
mock_create_async_engine):
+ """_configure_async_session() must not call the hook when
SQL_ALCHEMY_CONN_ASYNC is empty."""
+ from airflow import settings
+
+ with patch("airflow.settings.SQL_ALCHEMY_CONN_ASYNC", ""):
+ settings._configure_async_session()
+
+ mock_create_async_engine.assert_not_called()
+
+ @patch("airflow.settings.create_engine")
+ def test_default_create_metadata_engine_forwards_args(self,
mock_sa_create_engine):
+ """Default create_metadata_engine must forward all args to
sqlalchemy.create_engine."""
+ from airflow import settings
+
+ mock_sa_create_engine.return_value = MagicMock()
+ engine_args = {"pool_size": 5, "pool_recycle": 1800}
+ connect_args = {"timeout": 30}
+
+ settings.create_metadata_engine("sqlite://", engine_args=engine_args,
connect_args=connect_args)
+
+ mock_sa_create_engine.assert_called_once_with(
+ "sqlite://",
+ connect_args={"timeout": 30},
+ pool_size=5,
+ pool_recycle=1800,
+ future=True,
+ )
Review Comment:
```suggestion
assert mock_sa_create_engine.mock_calls == [
call(
"sqlite://",
connect_args={"timeout": 30},
pool_size=5,
pool_recycle=1800,
future=True,
)
]
```
--
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]
