lib/cloud_trace_unittest.py

# -*- coding: utf-8 -*-
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Tests for chromite.lib.cloud_trace."""
from __future__ import print_function

import os

from chromite.lib import cloud_trace
from chromite.lib import cros_test_lib


class SpanTest(cros_test_lib.MockTestCase):
  """Tests for the Span class."""

  def testToDict(self):
    """Tests that Span instances can be turned into dicts."""
    traceId = 'deadbeefdeadbeefdeadbeefdeadbeef'
    span = cloud_trace.Span(spanId='1234', traceId=traceId, name='Test span')
    self.assertEqual(span.ToDict(), {
        'labels': {},
        'spanId': '1234',
        'name': 'Test span'
    })


class SpanStackTest(cros_test_lib.MockTestCase):
  """Tests for the SpanStack class."""

  def setUp(self):
    self.log_span_mock = self.PatchObject(cloud_trace, 'LogSpan')

  def testLogSingleSpan(self):
    """Tests that SpanStack.Span logs a span and sends it."""
    stack = cloud_trace.SpanStack()
    context = stack.Span('foo')
    self.assertEqual(0, self.log_span_mock.call_count)
    with context:
      self.assertEqual(0, self.log_span_mock.call_count)
    self.assertEqual(1, self.log_span_mock.call_count)

  def testCallLogSpanAtCloseOfStack(self):
    """Test that LogSpans is called after each span is popped."""
    stack = cloud_trace.SpanStack()
    with stack.Span('foo'):
      self.assertEqual(0, self.log_span_mock.call_count)
      with stack.Span('bar'):
        self.assertEqual(0, self.log_span_mock.call_count)
        with stack.Span('zap'):
          self.assertEqual(0, self.log_span_mock.call_count)
        self.assertEqual(1, self.log_span_mock.call_count)
      self.assertEqual(2, self.log_span_mock.call_count)
    self.assertEqual(3, self.log_span_mock.call_count)

  def testSpannedDecorator(self):
    """Tests that @stack.Spanned() works."""
    stack = cloud_trace.SpanStack()
    @stack.Spanned('foo')
    def decorated():
      pass

    self.assertEqual(0, self.log_span_mock.call_count)
    decorated()
    self.assertEqual(1, self.log_span_mock.call_count)


class SpanLogTest(cros_test_lib.MockTestCase, cros_test_lib.TempDirTestCase):
  """Tests that spans can be logged correctly."""

  def setUp(self):
    self._old_log = cloud_trace.SPANS_LOG
    cloud_trace.SPANS_LOG = os.path.join(self.tempdir, '{pid}.json')

  def tearDown(self):
    cloud_trace.SPANS_LOG = self._old_log

  def testCanSendLog(self):
    """Tests that Spans are sent to a log."""
    stack = cloud_trace.SpanStack()
    with stack.Span('foo'):
      pass
    self.assertExists(cloud_trace.SPANS_LOG.format(pid=os.getpid()))


class CloudTraceContextTest(cros_test_lib.MockTestCase):
  """Tests the CLOUD_TRACE_CONTEXT environment variable functionality.."""
  def setUp(self):
    self.log_span_mock = self.PatchObject(cloud_trace, 'LogSpan')
    os_mock = self.PatchObject(cloud_trace, 'os')
    os_mock.environ = self.env = {}

  def testHeaderValue(self):
    """Tests that the header value makes sense."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    stack = cloud_trace.SpanStack(
        traceId=trace_id,
        parentSpanId='0',
        global_context='deadbeefdeadbeefdeadbeefdeadbeef/0;o=1')

    self.assertTrue(stack.enabled)
    self.assertEqual(stack.traceId, trace_id)
    self.assertEqual(stack.last_span_id, '0')

  def testCloudTraceContextPattern(self):
    """Tests that the regex matches an example context."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    global_context = '{trace_id}/0;o=1'.format(trace_id=trace_id)
    self.assertTrue(bool(
        cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_PATTERN.match(
            global_context)))

    # pylint: disable=protected-access
    result = cloud_trace.SpanStack._ParseCloudTraceContext(global_context)
    self.assertEqual(result, {
        'traceId': trace_id,
        'parentSpanId': '0',
        'options': '1'
    })

  def testInitReadsEnvironment(self):
    """Tests that SpanStack reads the enivornment on init."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    global_context = '{trace_id}/0;o=1'.format(trace_id=trace_id)
    old_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)
    try:
      self.env[cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV] = global_context
      stack = cloud_trace.SpanStack()
      self.assertEqual(stack.traceId, trace_id)
      self.assertEqual(stack.last_span_id, '0')
    finally:
      if old_env is not None:
        self.env[cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV] = old_env

  def testEnvironmentContextManager(self):
    """Tests that the environment context manager works."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    stack = cloud_trace.SpanStack(
        global_context='{trace_id}/0;o=1'.format(trace_id=trace_id))

    old_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)
    with stack.EnvironmentContext():
      new_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)
    after_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)

    self.assertEqual(old_env, after_env)
    # Note the lack of a /0; the /0 is optional.
    self.assertEqual(new_env, 'deadbeef12345678deadbeef12345678;o=1')

  def testSpanContextEnabled(self):
    """Tests that the span context manager updates the environment."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    stack = cloud_trace.SpanStack(
        traceId=trace_id,
        parentSpanId='0',
        global_context='deadbeefdeadbeefdeadbeefdeadbeef/0;o=1')

    self.assertTrue(stack.enabled)
    old_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)
    with stack.Span('foo') as span:
      new_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)
      self.assertTrue(new_env.startswith('deadbeef12345678deadbeef12345678/'))
      self.assertEqual(span.parentSpanId, '0')
      self.assertEqual(span.traceId, trace_id)
    after_env = self.env.get(cloud_trace.SpanStack.CLOUD_TRACE_CONTEXT_ENV)

    self.assertEqual(self.log_span_mock.call_count, 1)
    self.assertEqual(old_env, after_env)

  def testSpanContextDisabled(self):
    """Tests that o=1 in the global_context disables the spanstack."""
    trace_id = 'deadbeef12345678deadbeef12345678'
    stack = cloud_trace.SpanStack(
        traceId=trace_id,
        parentSpanId='0',
        global_context='deadbeefdeadbeefdeadbeefdeadbeef/0;o=0')

    self.assertFalse(stack.enabled)
    with stack.Span('foo') as span:
      self.assertEqual(span.parentSpanId, '0')
      self.assertEqual(span.traceId, trace_id)
    self.assertEqual(self.log_span_mock.call_count, 0)