Split code sugar and logic in Profiler class

Add new class Trace that can be used with "with" statement:

  with Trace(name, info):
     do_something

Remove "with" support from profiler class.

Add unit test and fix wsgi middleware according to this change.

Change-Id: Idf7c9bbf262b909da9865d8d396a1f62b8e4d48d

osprofiler/profiler.py

@@ -61,6 +61,19 @@ def stop(info=None):
         profiler.stop(info=info)
 
 
+class Trace(object):
+
+    def __init__(self, name, info=None):
+        self._name = name
+        self._info = info
+
+    def __enter__(self):
+        start(self._name, info=self._info)
+
+    def __exit__(self, etype, value, traceback):
+        stop()
+
+
 class Profiler(object):
 
     def __init__(self, base_id=None, parent_id=None, hmac_key=None):
@@ -70,22 +83,6 @@ class Profiler(object):
         self._trace_stack = collections.deque([base_id, parent_id or base_id])
         self._name = collections.deque()
 
-    def __call__(self, name, info=None):
-        """This method simplifies usage of profiler object as a guard
-        > profiler = Profiler(service='nova')
-        > with profiler('some long running code'):
-        >     do_some_stuff()
-        """
-        self._name.append(name)
-        self._info = info
-        return self
-
-    def __enter__(self):
-        self.start(self._name[-1], info=self._info)
-
-    def __exit__(self, etype, value, traceback):
-        self.stop()
-
     def get_base_id(self):
         return self._trace_stack[0]
 

osprofiler/web.py

@@ -84,9 +84,9 @@ class WsgiMiddleware(object):
             else:
                 trace_info = json.loads(trace_raw)
 
-                p = profiler.init(trace_info.get("base_id"),
-                                  trace_info.get("parent_id"),
-                                  self.hmac_key)
+                profiler.init(trace_info.get("base_id"),
+                              trace_info.get("parent_id"),
+                              self.hmac_key)
 
                 info = {
                     "request": {
@@ -98,7 +98,7 @@ class WsgiMiddleware(object):
                     }
                 }
 
-                with p(self.name, info=info):
+                with profiler.Trace(self.name, info=info):
                     return request.get_response(self.application)
 
         return request.get_response(self.application)

tests/test_profiler.py

@@ -123,18 +123,18 @@ class ProfilerTestCase(test.TestCase):
         self.assertEqual(len(prof._name), 0)
         self.assertEqual(prof._trace_stack, collections.deque(["1", "2"]))
 
-    def test_profiler_with_statement(self):
-        prof = profiler.Profiler(base_id="1", parent_id="2")
-        prof.start = mock.MagicMock()
-        prof.stop = mock.MagicMock()
 
-        with prof("name1", info="test"):
-            prof.start.assert_called_once_with("name1", info="test")
-            prof.start.reset_mock()
-            self.assertFalse(prof.stop.called)
-            with prof("name2", info="test2"):
-                prof.start.assert_called_once_with("name2", info="test2")
-                self.assertFalse(prof.stop.called)
-            prof.stop.assert_called_once_with()
-            prof.stop.reset_mock()
-        prof.stop.assert_called_once_with()
+class TraceTestCase(test.TestCase):
+
+    @mock.patch("osprofiler.profiler.stop")
+    @mock.patch("osprofiler.profiler.start")
+    def test_trace(self, mock_start, mock_stop):
+
+        with profiler.Trace("a", info="a1"):
+            mock_start.assert_called_once_with("a", info="a1")
+            mock_start.reset_mock()
+            with profiler.Trace("b", info="b1"):
+                mock_start.assert_called_once_with("b", info="b1")
+            mock_stop.assert_called_once_with()
+            mock_stop.reset_mock()
+        mock_stop.assert_called_once_with()

tests/test_web.py

@@ -16,7 +16,6 @@
 import json
 import mock
 
-from osprofiler import profiler
 from osprofiler import utils
 from osprofiler import web
 
@@ -74,7 +73,9 @@ class WebMiddlewareTestCase(test.TestCase):
         self.assertEqual("yeah!", middleware(request))
         request.get_response.assert_called_once_with("app")
 
-    def test_wsgi_middleware(self):
+    @mock.patch("osprofiler.web.profiler.Trace")
+    @mock.patch("osprofiler.web.profiler.init")
+    def test_wsgi_middleware(self, mock_profiler_init, mock_profiler_trace):
         request = mock.MagicMock()
         request.get_response.return_value = "yeah!"
         request.url = "someurl"
@@ -90,24 +91,17 @@ class WebMiddlewareTestCase(test.TestCase):
             "b": "2",
             "X-Trace-Info": utils.binary_encode(json.dumps(trace_info))
         }
-        p = profiler.init()
-        p.start = mock.MagicMock()
-        p.stop = mock.MagicMock()
 
-        with mock.patch("osprofiler.web.profiler.init") as mock_profiler_init:
-            mock_profiler_init.return_value = p
-
-            middleware = web.WsgiMiddleware("app", enabled=True)
-            self.assertEqual("yeah!", middleware(request))
-            mock_profiler_init.assert_called_once_with("1", "2", None)
-            expected_info = {
-                "request": {
-                    "host_url": request.host_url,
-                    "path": request.path,
-                    "query": request.query_string,
-                    "method": request.method,
-                    "scheme": request.scheme
-                }
+        middleware = web.WsgiMiddleware("app", enabled=True)
+        self.assertEqual("yeah!", middleware(request))
+        mock_profiler_init.assert_called_once_with("1", "2", None)
+        expected_info = {
+            "request": {
+                "host_url": request.host_url,
+                "path": request.path,
+                "query": request.query_string,
+                "method": request.method,
+                "scheme": request.scheme
             }
-            p.start.assert_called_once_with("wsgi", info=expected_info)
-            p.stop.assert_called_once_with()
+        }
+        mock_profiler_trace.assert_called_once_with("wsgi", info=expected_info)