From 3c8f7da7a4744b02062978c1af37ba372ffc3360 Mon Sep 17 00:00:00 2001
From: Boris Pavlovic <boris@pavlovic.me>
Date: Sun, 23 Feb 2014 23:24:58 +0400
Subject: [PATCH] Remove from sqlachemy.after_execute notifcation resutls and
 add UTs

Mostly this patch just covers by UTs sqlalcehmy module
---
 osprofiler/sqlalchemy.py | 21 +++++++-------
 tests/test_sqlalchemy.py | 60 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 11 deletions(-)
 create mode 100644 tests/test_sqlalchemy.py

diff --git a/osprofiler/sqlalchemy.py b/osprofiler/sqlalchemy.py
index 4ddf580..1a8d793 100644
--- a/osprofiler/sqlalchemy.py
+++ b/osprofiler/sqlalchemy.py
@@ -17,29 +17,28 @@ from osprofiler import profiler
 
 
 def before_execute(name):
-    """Add listener that will send trace info before sql executed."""
+    """Add listener that will send trace info before query is executed."""
+
     def handler(conn, clauseelement, multiparams, params):
-        p = profiler.get_profiler()
-        if p:
-            info = {"db.statement": str(clauseelement),
-                    "db.multiparams": str(multiparams),
-                    "db.params": str(params)}
-            p.start(name, info=info)
+        info = {"db.statement": str(clauseelement),
+                "db.multiparams": str(multiparams),
+                "db.params": str(params)}
+        profiler.start(name, info=info)
 
     return handler
 
 
 def after_execute():
-    """Add listener that will send trace info after sql executed."""
+    """Add listener that will send trace info after query is executed."""
+
     def handler(conn, clauseelement, multiparams, params, result):
-        p = profiler.get_profiler()
-        if p:
-            p.stop(info={"db.result": str(result)})
+        profiler.stop(info=None)
 
     return handler
 
 
 def add_tracing(sqlalchemy, engine, name):
     """Add tracing to all sqlalchemy calls."""
+
     sqlalchemy.event.listen(engine, 'before_execute', before_execute(name))
     sqlalchemy.event.listen(engine, 'after_execute', after_execute())
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
new file mode 100644
index 0000000..66e570c
--- /dev/null
+++ b/tests/test_sqlalchemy.py
@@ -0,0 +1,60 @@
+# Copyright 2014 Mirantis Inc.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import mock
+
+from osprofiler import sqlalchemy
+
+from tests import test
+
+
+class SqlalchemyTracingTestCase(test.TestCase):
+
+    @mock.patch("osprofiler.sqlalchemy.profiler")
+    def test_before_execute(self, mock_profiler):
+        handler = sqlalchemy.before_execute("sql")
+
+        handler(mock.MagicMock(), 1, 2, 3)
+        expected_info = {
+            "db.statement": "1",
+            "db.multiparams": "2",
+            "db.params": "3"
+        }
+        mock_profiler.start.assert_called_once_with("sql", info=expected_info)
+
+    @mock.patch("osprofiler.sqlalchemy.profiler")
+    def test_after_execute(self, mock_profiler):
+        handler = sqlalchemy.after_execute()
+        handler(mock.MagicMock(), 1, 2, 3, mock.MagicMock())
+        mock_profiler.stop.assert_called_once_with(info=None)
+
+    @mock.patch("osprofiler.sqlalchemy.before_execute")
+    @mock.patch("osprofiler.sqlalchemy.after_execute")
+    def test_add_tracing(self, mock_after_exc, mock_before_exc):
+        sa = mock.MagicMock()
+        engine = mock.MagicMock()
+
+        mock_before_exc.return_value = "before"
+        mock_after_exc.return_value = "after"
+
+        sqlalchemy.add_tracing(sa, engine, "sql")
+
+        mock_before_exc.assert_called_once_with("sql")
+        mock_after_exc.assert_called_once_with()
+        expected_calls = [
+            mock.call(engine, "before_execute", "before"),
+            mock.call(engine, "after_execute", "after")
+        ]
+        self.assertEqual(sa.event.listen.call_args_list, expected_calls)