diff options
Diffstat (limited to 'osprofiler/tests/test_sqlalchemy.py')
| -rw-r--r-- | osprofiler/tests/test_sqlalchemy.py | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/osprofiler/tests/test_sqlalchemy.py b/osprofiler/tests/test_sqlalchemy.py index 1494da6..83ada93 100644 --- a/osprofiler/tests/test_sqlalchemy.py +++ b/osprofiler/tests/test_sqlalchemy.py @@ -13,6 +13,7 @@ # License for the specific language governing permissions and limitations # under the License. +import contextlib import mock from osprofiler import sqlalchemy @@ -56,6 +57,36 @@ class SqlalchemyTracingTestCase(test.TestCase): @mock.patch("osprofiler.sqlalchemy._before_cursor_execute") @mock.patch("osprofiler.sqlalchemy._after_cursor_execute") + def test_wrap_session(self, mock_after_exc, mock_before_exc): + sa = mock.MagicMock() + + @contextlib.contextmanager + def _session(): + session = mock.MagicMock() + # current engine object stored within the session + session.bind = mock.MagicMock() + session.bind.traced = None + yield session + + mock_before_exc.return_value = "before" + mock_after_exc.return_value = "after" + + session = sqlalchemy.wrap_session(sa, _session()) + + with session as sess: + pass + + mock_before_exc.assert_called_once_with("db") + mock_after_exc.assert_called_once_with() + expected_calls = [ + mock.call(sess.bind, "before_cursor_execute", "before"), + mock.call(sess.bind, "after_cursor_execute", "after") + ] + + self.assertEqual(sa.event.listen.call_args_list, expected_calls) + + @mock.patch("osprofiler.sqlalchemy._before_cursor_execute") + @mock.patch("osprofiler.sqlalchemy._after_cursor_execute") def test_disable_and_enable(self, mock_after_exc, mock_before_exc): sqlalchemy.disable() |