xref: /aosp_15_r20/build/make/tools/tool_event_logger/tool_event_logger_test.py (revision 9e94795a3d4ef5c1d47486f9a02bb378756cea8a) 
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_HOST_NAME = 'test_host_name'
29TEST_TOOL_TAG = 'test_tool'
30TEST_SOURCE_ROOT = 'test_source_root'
31TEST_PLATFORM_VERSION = 'test_platform_version'
32TEST_PYTHON_VERSION = 'test_python_version'
33TEST_EVENT_TIMESTAMP = datetime.datetime.now()
34
35
36class ToolEventLoggerTest(unittest.TestCase):
37
38  def setUp(self):
39    super().setUp()
40    self.clearcut_client = FakeClearcutClient()
41    self.logger = tool_event_logger.ToolEventLogger(
42        TEST_TOOL_TAG,
43        TEST_INVOCATION_ID,
44        TEST_USER_NAME,
45        TEST_HOST_NAME,
46        TEST_SOURCE_ROOT,
47        TEST_PLATFORM_VERSION,
48        TEST_PYTHON_VERSION,
49        client=self.clearcut_client,
50    )
51
52  def test_log_event_timestamp(self):
53    with self.logger:
54      self.logger.log_invocation_started(
55          datetime.datetime.fromtimestamp(100.101), 'test_command'
56      )
57
58    self.assertEqual(
59        self.clearcut_client.get_last_sent_event().event_time_ms, 100101
60    )
61
62  def test_log_event_basic_information(self):
63    with self.logger:
64      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
65
66    sent_event = self.clearcut_client.get_last_sent_event()
67    log_event = tool_event_pb2.ToolEvent.FromString(sent_event.source_extension)
68    self.assertEqual(log_event.invocation_id, TEST_INVOCATION_ID)
69    self.assertEqual(log_event.user_name, TEST_USER_NAME)
70    self.assertEqual(log_event.host_name, TEST_HOST_NAME)
71    self.assertEqual(log_event.tool_tag, TEST_TOOL_TAG)
72    self.assertEqual(log_event.source_root, TEST_SOURCE_ROOT)
73
74  def test_log_invocation_started(self):
75    expected_invocation_started = tool_event_pb2.ToolEvent.InvocationStarted(
76        command_args='test_command',
77        os=TEST_PLATFORM_VERSION + ':' + TEST_PYTHON_VERSION,
78    )
79
80    with self.logger:
81      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
82
83    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
84    sent_event = self.clearcut_client.get_last_sent_event()
85    self.assertEqual(
86        expected_invocation_started,
87        tool_event_pb2.ToolEvent.FromString(
88            sent_event.source_extension
89        ).invocation_started,
90    )
91
92  def test_log_invocation_stopped(self):
93    expected_invocation_stopped = tool_event_pb2.ToolEvent.InvocationStopped(
94        exit_code=0,
95        exit_log='exit_log',
96    )
97
98    with self.logger:
99      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')
100
101    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
102    sent_event = self.clearcut_client.get_last_sent_event()
103    self.assertEqual(
104        expected_invocation_stopped,
105        tool_event_pb2.ToolEvent.FromString(
106            sent_event.source_extension
107        ).invocation_stopped,
108    )
109
110  def test_log_multiple_events(self):
111    with self.logger:
112      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
113      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')
114
115    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 2)
116
117
118class MainTest(unittest.TestCase):
119
120  REQUIRED_ARGS = [
121      '',
122      '--tool_tag',
123      'test_tool',
124      '--start_timestamp',
125      '1',
126      '--end_timestamp',
127      '2',
128      '--exit_code',
129      '0',
130  ]
131
132  def test_log_and_exit_with_missing_required_args(self):
133    with self.assertLogs() as logs:
134      with self.assertRaises(SystemExit) as ex:
135        tool_event_logger.main(['', '--tool_tag', 'test_tool'])
136
137    with self.subTest('Verify exception code'):
138      self.assertEqual(ex.exception.code, 2)
139
140    with self.subTest('Verify log messages'):
141      self.assertIn(
142          'the following arguments are required',
143          '\n'.join(logs.output),
144      )
145
146  def test_log_and_exit_with_invalid_args(self):
147    with self.assertLogs() as logs:
148      with self.assertRaises(SystemExit) as ex:
149        tool_event_logger.main(['', '--start_timestamp', 'test'])
150
151    with self.subTest('Verify exception code'):
152      self.assertEqual(ex.exception.code, 2)
153
154    with self.subTest('Verify log messages'):
155      self.assertIn(
156          '--start_timestamp: invalid',
157          '\n'.join(logs.output),
158      )
159
160  def test_log_and_exit_with_dry_run(self):
161    with self.assertLogs(level=logging.DEBUG) as logs:
162      tool_event_logger.main(self.REQUIRED_ARGS + ['--dry_run'])
163
164    with self.subTest('Verify log messages'):
165      self.assertIn('dry run', '\n'.join(logs.output))
166
167  @mock.patch.object(clearcut_client, 'Clearcut')
168  def test_log_and_exit_with_unexpected_exception(self, mock_cc):
169    mock_cc.return_value = FakeClearcutClient(raise_log_exception=True)
170
171    with self.assertLogs() as logs:
172      with self.assertRaises(Exception) as ex:
173        tool_event_logger.main(self.REQUIRED_ARGS)
174
175    with self.subTest('Verify log messages'):
176      self.assertIn('unexpected error', '\n'.join(logs.output))
177
178  @mock.patch.object(clearcut_client, 'Clearcut')
179  def test_success(self, mock_cc):
180    mock_clear_cut_client = FakeClearcutClient()
181    mock_cc.return_value = mock_clear_cut_client
182
183    tool_event_logger.main(self.REQUIRED_ARGS)
184
185    self.assertEqual(mock_clear_cut_client.get_number_of_sent_events(), 2)
186
187
188class FakeClearcutClient:
189
190  def __init__(self, raise_log_exception=False):
191    self.pending_log_events = []
192    self.sent_log_events = []
193    self.raise_log_exception = raise_log_exception
194
195  def log(self, log_event):
196    if self.raise_log_exception:
197      raise Exception('unknown exception')
198    self.pending_log_events.append(log_event)
199
200  def flush_events(self):
201    self.sent_log_events.extend(self.pending_log_events)
202    self.pending_log_events.clear()
203
204  def get_number_of_sent_events(self):
205    return len(self.sent_log_events)
206
207  def get_last_sent_event(self):
208    return self.sent_log_events[-1]
209
210
211if __name__ == '__main__':
212  unittest.main()
213