• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2024, The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Unittests for ToolEventLogger."""
16
17import datetime
18import logging
19import unittest
20from unittest import mock
21
22from atest.metrics import clearcut_client
23from proto import tool_event_pb2
24from tool_event_logger import tool_event_logger
25
26TEST_INVOCATION_ID = 'test_invocation_id'
27TEST_USER_NAME = 'test_user'
28TEST_TOOL_TAG = 'test_tool'
29TEST_SOURCE_ROOT = 'test_source_root'
30TEST_PLATFORM_VERSION = 'test_platform_version'
31TEST_PYTHON_VERSION = 'test_python_version'
32TEST_EVENT_TIMESTAMP = datetime.datetime.now()
33
34
35class ToolEventLoggerTest(unittest.TestCase):
36
37  def setUp(self):
38    super().setUp()
39    self.clearcut_client = FakeClearcutClient()
40    self.logger = tool_event_logger.ToolEventLogger(
41        TEST_TOOL_TAG,
42        TEST_INVOCATION_ID,
43        TEST_USER_NAME,
44        TEST_SOURCE_ROOT,
45        TEST_PLATFORM_VERSION,
46        TEST_PYTHON_VERSION,
47        client=self.clearcut_client,
48    )
49
50  def test_log_event_timestamp(self):
51    with self.logger:
52      self.logger.log_invocation_started(
53          datetime.datetime.fromtimestamp(100.101), 'test_command'
54      )
55
56    self.assertEqual(
57        self.clearcut_client.get_last_sent_event().event_time_ms, 100101
58    )
59
60  def test_log_event_basic_information(self):
61    with self.logger:
62      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
63
64    sent_event = self.clearcut_client.get_last_sent_event()
65    log_event = tool_event_pb2.ToolEvent.FromString(sent_event.source_extension)
66    self.assertEqual(log_event.invocation_id, TEST_INVOCATION_ID)
67    self.assertEqual(log_event.user_name, TEST_USER_NAME)
68    self.assertEqual(log_event.tool_tag, TEST_TOOL_TAG)
69    self.assertEqual(log_event.source_root, TEST_SOURCE_ROOT)
70
71  def test_log_invocation_started(self):
72    expected_invocation_started = tool_event_pb2.ToolEvent.InvocationStarted(
73        command_args='test_command',
74        os=TEST_PLATFORM_VERSION + ':' + TEST_PYTHON_VERSION,
75    )
76
77    with self.logger:
78      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
79
80    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
81    sent_event = self.clearcut_client.get_last_sent_event()
82    self.assertEqual(
83        expected_invocation_started,
84        tool_event_pb2.ToolEvent.FromString(
85            sent_event.source_extension
86        ).invocation_started,
87    )
88
89  def test_log_invocation_stopped(self):
90    expected_invocation_stopped = tool_event_pb2.ToolEvent.InvocationStopped(
91        exit_code=0,
92        exit_log='exit_log',
93    )
94
95    with self.logger:
96      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')
97
98    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
99    sent_event = self.clearcut_client.get_last_sent_event()
100    self.assertEqual(
101        expected_invocation_stopped,
102        tool_event_pb2.ToolEvent.FromString(
103            sent_event.source_extension
104        ).invocation_stopped,
105    )
106
107  def test_log_multiple_events(self):
108    with self.logger:
109      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
110      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')
111
112    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 2)
113
114
115class MainTest(unittest.TestCase):
116
117  REQUIRED_ARGS = [
118      '',
119      '--tool_tag',
120      'test_tool',
121      '--start_timestamp',
122      '1',
123      '--end_timestamp',
124      '2',
125      '--exit_code',
126      '0',
127  ]
128
129  def test_log_and_exit_with_missing_required_args(self):
130    with self.assertLogs() as logs:
131      with self.assertRaises(SystemExit) as ex:
132        tool_event_logger.main(['', '--tool_tag', 'test_tool'])
133
134    with self.subTest('Verify exception code'):
135      self.assertEqual(ex.exception.code, 2)
136
137    with self.subTest('Verify log messages'):
138      self.assertIn(
139          'the following arguments are required',
140          '\n'.join(logs.output),
141      )
142
143  def test_log_and_exit_with_invalid_args(self):
144    with self.assertLogs() as logs:
145      with self.assertRaises(SystemExit) as ex:
146        tool_event_logger.main(['', '--start_timestamp', 'test'])
147
148    with self.subTest('Verify exception code'):
149      self.assertEqual(ex.exception.code, 2)
150
151    with self.subTest('Verify log messages'):
152      self.assertIn(
153          '--start_timestamp: invalid',
154          '\n'.join(logs.output),
155      )
156
157  def test_log_and_exit_with_dry_run(self):
158    with self.assertLogs(level=logging.DEBUG) as logs:
159      tool_event_logger.main(self.REQUIRED_ARGS + ['--dry_run'])
160
161    with self.subTest('Verify log messages'):
162      self.assertIn('dry run', '\n'.join(logs.output))
163
164  @mock.patch.object(clearcut_client, 'Clearcut')
165  def test_log_and_exit_with_unexpected_exception(self, mock_cc):
166    mock_cc.return_value = FakeClearcutClient(raise_log_exception=True)
167
168    with self.assertLogs() as logs:
169      with self.assertRaises(Exception) as ex:
170        tool_event_logger.main(self.REQUIRED_ARGS)
171
172    with self.subTest('Verify log messages'):
173      self.assertIn('unexpected error', '\n'.join(logs.output))
174
175  @mock.patch.object(clearcut_client, 'Clearcut')
176  def test_success(self, mock_cc):
177    mock_clear_cut_client = FakeClearcutClient()
178    mock_cc.return_value = mock_clear_cut_client
179
180    tool_event_logger.main(self.REQUIRED_ARGS)
181
182    self.assertEqual(mock_clear_cut_client.get_number_of_sent_events(), 2)
183
184
185class FakeClearcutClient:
186
187  def __init__(self, raise_log_exception=False):
188    self.pending_log_events = []
189    self.sent_log_events = []
190    self.raise_log_exception = raise_log_exception
191
192  def log(self, log_event):
193    if self.raise_log_exception:
194      raise Exception('unknown exception')
195    self.pending_log_events.append(log_event)
196
197  def flush_events(self):
198    self.sent_log_events.extend(self.pending_log_events)
199    self.pending_log_events.clear()
200
201  def get_number_of_sent_events(self):
202    return len(self.sent_log_events)
203
204  def get_last_sent_event(self):
205    return self.sent_log_events[-1]
206
207
208if __name__ == '__main__':
209  unittest.main()
210