|
17 | 17 | # under the License. |
18 | 18 | from __future__ import annotations |
19 | 19 |
|
| 20 | +import json |
20 | 21 | from unittest import mock |
21 | 22 | from unittest.mock import Mock |
22 | 23 |
|
@@ -135,3 +136,33 @@ def test_has_access( |
135 | 136 | if len(auth_manager_methods) > 1 and not expected: |
136 | 137 | for method_name in auth_manager_methods: |
137 | 138 | getattr(auth_manager, method_name).assert_called() |
| 139 | + |
| 140 | + @mock.patch("airflow.utils.session.create_session") |
| 141 | + @mock.patch("airflow.www.security_manager.get_auth_manager") |
| 142 | + def test_manager_does_not_create_extra_db_sessions( |
| 143 | + self, |
| 144 | + _, |
| 145 | + mock_create_session, |
| 146 | + security_manager, |
| 147 | + ): |
| 148 | + """ |
| 149 | + Test that the Security Manager doesn't create extra DB sessions and |
| 150 | + instead uses the session already available through the appbuilder |
| 151 | + object that is attached to it. |
| 152 | + """ |
| 153 | + with mock.patch.object(security_manager.appbuilder, "session") as mock_appbuilder_session: |
| 154 | + action_name = ACTION_CAN_READ |
| 155 | + resource_pk = "PK" |
| 156 | + user = Mock() |
| 157 | + for func in security_manager._auth_manager_is_authorized_map.values(): |
| 158 | + try: |
| 159 | + func(action_name, resource_pk, user) |
| 160 | + except json.JSONDecodeError: |
| 161 | + # The resource-retrieving function expects a "composite" |
| 162 | + # PK as a JSON string. Provide a mocked one. |
| 163 | + func(action_name, "[1, 1, 1, 1]", user) |
| 164 | + mock_create_session.assert_not_called() |
| 165 | + |
| 166 | + # The Security Manager's `appbuilder.session` object should have been |
| 167 | + # put to use by many of the functions tested above. |
| 168 | + assert len(mock_appbuilder_session.method_calls) > 0 |
0 commit comments