Change to use itertools to traverse spans

Author: liyanhui1228

trace/google/cloud/trace/trace.py

@@ -92,8 +92,9 @@ def send(self):
         """
         spans_list = []
         for root_span in self.spans:
-            span_tree = _traverse_span_tree(root_span)
-            spans_list.extend(span_tree)
+            span_tree = list(iter(root_span))
+            span_tree_json = [format_span_json(span) for span in span_tree]
+            spans_list.extend(span_tree_json)
 
         if len(spans_list) == 0:
             return
@@ -113,34 +114,7 @@ def send(self):
             traces=traces,
             options=None)
 
-        return traces
-
-
-def _traverse_span_tree(root_span):
-    """Helper to traverse the span tree in level order.
-
-    :rtype: :class:`~google.cloud.trace.trace_span.TraceSpan`
-    :param root_span: The root span in a span tree.
-
-    :rtype: list
-    :returns: A list of all the spans in a span tree.
-    """
-    span_list = []
-
-    if root_span is None:
-        return span_list
-
-    span_queue = []
-    span_queue.append(root_span)
-
-    while span_queue:
-        cur_span = span_queue.pop(0)
-        span_list.append(format_span_json(cur_span))
-
-        for child_span in cur_span.children:
-            span_queue.append(child_span)
-
-    return span_list
+        return spans_list
 
 
 def generate_trace_id():

trace/google/cloud/trace/trace_span.py

@@ -17,6 +17,7 @@
 from datetime import datetime
 from google.cloud.gapic.trace.v1.enums import TraceSpan as Enum
 
+from itertools import chain
 import random
 
 
@@ -112,6 +113,11 @@ def set_end_time(self):
         """Set the end time for a span."""
         self.end_time = datetime.utcnow().isoformat() + 'Z'
 
+    def __iter__(self):
+        for span in chain(*(map(iter, self.children))):
+            yield span
+        yield self
+
     def __enter__(self):
         self.set_start_time()
         return self

trace/tests/unit/test_trace.py

@@ -65,24 +65,9 @@ def test_start(self):
 
         self.assertEqual(trace.spans, [])
 
-    def test_finish_with_empty_span(self):
-
-        def patch_traces(traces, project_id=None, options=None):
-            _patch_traces_called_with = (traces, project_id, options)
-            return _patch_traces_called_with
-
-        client = mock.Mock(project=self.project, spec=['project'])
-        client.patch_traces = patch_traces
-        trace = self._make_one(client=client)
-
-        with trace:
-            trace.spans = [None]
-            self.assertEqual(trace.spans, [None])
-
-        self.assertEqual(trace.spans, [])
-
     def test_finish_with_valid_span(self):
         from google.cloud.gapic.trace.v1.enums import TraceSpan as Enum
+        from google.cloud.trace.trace_span import TraceSpan
 
         def patch_traces(traces, project_id=None, options=None):
             _patch_traces_called_with = (traces, project_id, options)
@@ -98,22 +83,16 @@ def patch_traces(traces, project_id=None, options=None):
         start_time = '2017-06-25'
         end_time = '2017-06-26'
 
-        span = mock.Mock(name=span_name,
-                         kind=kind,
-                         parent_span_id=None,
-                         span_id=span_id,
-                         start_time=start_time,
-                         end_time=end_time,
-                         labels=None,
-                         children=[],
-                         spec=['name',
-                               'kind',
-                               'parent_span_id',
-                               'span_id',
-                               'start_time',
-                               'end_time',
-                               'labels',
-                               'children'])
+        span = mock.Mock(spec=TraceSpan)
+        span.name = span_name
+        span.kind = kind
+        span.parent_span_id = None
+        span.span_id = span_id
+        span.start_time = start_time
+        span.end_time = end_time
+        span.labels = None
+        span.children = []
+        span.__iter__ = mock.Mock(return_value=iter([span]))
 
         with trace:
             trace.spans = [span]
@@ -153,6 +132,7 @@ def test_send_without_spans(self):
     def test_send_with_spans(self):
         from google.cloud.gapic.trace.v1.enums import TraceSpan as Enum
         from google.cloud.trace.client import Client
+        from google.cloud.trace.trace_span import TraceSpan
 
         client = mock.Mock(spec=Client)
         client.project = self.project
@@ -170,38 +150,28 @@ def test_send_with_spans(self):
             '/component': 'HTTP load balancer',
         }
 
-        child_span = mock.Mock(name=child_span_name,
-                               kind=kind,
-                               parent_span_id=root_span_id,
-                               span_id=child_span_id,
-                               start_time=start_time,
-                               end_time=end_time,
-                               labels=labels,
-                               children=[],
-                               spec=['name',
-                                     'kind',
-                                     'parent_span_id',
-                                     'span_id',
-                                     'start_time',
-                                     'end_time',
-                                     'labels',
-                                     'children'])
-        root_span = mock.Mock(name=root_span_name,
-                              kind=kind,
-                              parent_span_id=None,
-                              span_id=root_span_id,
-                              start_time=start_time,
-                              end_time=end_time,
-                              labels=None,
-                              children=[child_span],
-                              spec=['name',
-                                    'kind',
-                                    'parent_span_id',
-                                    'span_id',
-                                    'start_time',
-                                    'end_time',
-                                    'labels',
-                                    'children'])
+        child_span = mock.Mock(spec=TraceSpan)
+        child_span.name = child_span_name
+        child_span.kind = kind
+        child_span.parent_span_id = root_span_id
+        child_span.span_id = child_span_id
+        child_span.start_time = start_time
+        child_span.end_time = end_time
+        child_span.labels = labels
+        child_span.children = []
+        child_span.__iter__ = mock.Mock(return_value=iter([child_span]))
+
+        root_span = mock.Mock(spec=TraceSpan)
+        root_span.name = root_span_name
+        root_span.kind = kind
+        root_span.parent_span_id = None
+        root_span.span_id = root_span_id
+        root_span.start_time = start_time
+        root_span.end_time = end_time
+        root_span.labels = None
+        root_span.children = []
+        root_span.__iter__ = mock.Mock(
+            return_value=iter([root_span, child_span]))
 
         child_span_json = {
             'name': child_span.name,

trace/tests/unit/test_trace_span.py

@@ -131,3 +131,23 @@ def test_set_end_time(self):
 
         span.set_end_time()
         self.assertIsNotNone(span.end_time)
+
+    def test___iter__(self):
+        root_span_name = 'root_span_name'
+        child1_span_name = 'child1_span_name'
+        child2_span_name = 'child2_span_name'
+        child1_child1_span_name = 'child1_child1_span_name'
+
+        root_span = self._make_one(root_span_name)
+        child1_span = self._make_one(child1_span_name)
+        child2_span = self._make_one(child2_span_name)
+        child1_child1_span = self._make_one(child1_child1_span_name)
+
+        child1_span._child_spans.append(child1_child1_span)
+        root_span._child_spans.extend([child1_span, child2_span])
+
+        span_iter_list = list(iter(root_span))
+
+        self.assertEqual(
+            span_iter_list,
+            [child1_child1_span, child1_span, child2_span, root_span])