| # -*- coding: utf-8 -*- |
| # Copyright 2016 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. |
| |
| """Unittests for export_to_cloud_trace.""" |
| |
| from __future__ import print_function |
| |
| import contextlib |
| import sys |
| import threading |
| |
| import mock |
| |
| from chromite.lib import cros_test_lib |
| from chromite.scripts import export_to_cloud_trace |
| |
| |
| assert sys.version_info >= (3, 6), 'This module requires Python 3.6+' |
| |
| |
| _SPAN = """{ |
| "name": "foo", |
| "startTime": "2017-08-23T23:28:46.484326Z", |
| "endTime": "2017-08-23T23:29:46.484326Z" |
| }""" |
| |
| # pylint: disable=protected-access |
| class ExportToCloudTraceTest(cros_test_lib.MockTempDirTestCase): |
| """Test that various functions in export_to_cloud_trace work.""" |
| |
| def setUp(self): |
| # Put a limit on how many times "time" can be called; this prevents the |
| # thread from spinning forever if there is a bug. |
| self.PatchObject( |
| export_to_cloud_trace, 'time', |
| mock.Mock(time=mock.Mock(side_effect=list(range(42))))) |
| |
| self.PatchObject(export_to_cloud_trace, '_ReadAndDeletePreexisting') |
| |
| @contextlib.contextmanager |
| def _SendSpansThread(self): |
| client = mock.Mock() |
| thread = threading.Thread( |
| target=lambda: export_to_cloud_trace._WatchAndSendSpans( |
| 'example-project', client)) |
| thread.start() |
| yield thread, client |
| thread.join() |
| |
| def _PatchDirWatcher(self, spans): |
| dir_watcher = mock.MagicMock(Batches=lambda: iter(spans)) |
| dir_watcher.__enter__ = mock.Mock(return_value=dir_watcher) |
| self.PatchObject( |
| export_to_cloud_trace, '_DirWatcher', mock.Mock( |
| return_value=dir_watcher)) |
| |
| def testSendSingleEvent(self): |
| """Test that the script waits for a larger batch before sending.""" |
| self.PatchObject(export_to_cloud_trace, 'MIN_BATCH_SIZE', 2) |
| self._PatchDirWatcher(spans=[ |
| [_SPAN] |
| ]) |
| with self._SendSpansThread() as (_, client): |
| pass |
| |
| # Even though MIN_BATCH_SIZE is 2, we will send a smaller batch if we run |
| # out of time. |
| self.assertEqual(client.projects().patchTraces.call_count, 1) |
| |
| def testSendBatches(self): |
| """Test that the script waits for a larger batch before sending.""" |
| self.PatchObject(export_to_cloud_trace, 'MIN_BATCH_SIZE', 2) |
| error_log = self.PatchObject(export_to_cloud_trace.log, 'error') |
| self._PatchDirWatcher(spans=[ |
| [_SPAN, _SPAN], [_SPAN, _SPAN, _SPAN] |
| ]) |
| with self._SendSpansThread() as (_, client): |
| pass |
| |
| # MIN_BATCH_SIZE=2, but the second batch has 3 elements, and gets sent as |
| # one batch. |
| self.assertEqual(client.projects().patchTraces.call_count, 2) |
| # We log an error on metrics emission when the span does not have required |
| # fields (endTime, startTime and name). |
| self.assertEqual(0, error_log.call_count) |