Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit 836b9e1b authored by Zhuoyao Zhang's avatar Zhuoyao Zhang Committed by Gerrit Code Review
Browse files

Merge "Add a Clearcut tool event logger" into main

parents 0da2fee8 7b11b710
Loading
Loading
Loading
Loading
+67 −0
Original line number Diff line number Diff line
// Copyright 2024 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Set of error prone rules to ensure code quality
// PackageLocation check requires the androidCompatible=false otherwise it does not do anything.

package {
    default_applicable_licenses: ["Android-Apache-2.0"],
    default_team: "trendy_team_adte",
}

python_library_host {
    name: "tool_event_proto",
    srcs: [
        "proto/tool_event.proto",
    ],
    proto: {
        canonical_path_from_root: false,
    },
}

python_binary_host {
    name: "tool_event_logger",
    pkg_path: "tool_event_logger",
    srcs: [
        "tool_event_logger.py",
    ],
    libs: [
        "asuite_cc_client",
        "tool_event_proto",
    ],
    main: "tool_event_logger.py",
}

python_test_host {
    name: "tool_event_logger_test",
    main: "tool_event_logger_test.py",
    pkg_path: "tool_event_logger",
    srcs: [
        "tool_event_logger.py",
        "tool_event_logger_test.py",
    ],
    test_options: {
        unit_test: true,
    },
    libs: [
        "asuite_cc_client",
        "tool_event_proto",
    ],
    version: {
        py3: {
            embedded_launcher: true,
            enabled: true,
        },
    },
}
+4 −0
Original line number Diff line number Diff line
include platform/tools/asuite:/OWNERS

zhuoyao@google.com
hzalek@google.com
 No newline at end of file
+35 −0
Original line number Diff line number Diff line
syntax = "proto3";

package tools.asuite.tool_event_logger;

message ToolEvent {
  // Occurs immediately upon execution of the tool.
  message InvocationStarted {
    string command_args = 1;
    string cwd = 2;
    string os = 3;
  }

  // Occurs when tool exits for any reason.
  message InvocationStopped {
    int32 exit_code = 2;
    string exit_log = 3;
  }

  // ------------------------
  // FIELDS FOR ToolEvent
  // ------------------------
  // Random string generated to identify the invocation.
  string invocation_id = 1;
  // Internal user name.
  string user_name = 2;
  // The root of Android source.
  string source_root = 3;
  // Name of the tool used.
  string tool_tag = 6;

  oneof event {
    InvocationStarted invocation_started = 4;
    InvocationStopped invocation_stopped = 5;
  }
}
+229 −0
Original line number Diff line number Diff line
# Copyright 2024, The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import datetime
import getpass
import logging
import os
import platform
import sys
import tempfile
import uuid

from atest.metrics import clearcut_client
from atest.proto import clientanalytics_pb2
from proto import tool_event_pb2

LOG_SOURCE = 2395


class ToolEventLogger:
  """Logs tool events to Sawmill through Clearcut."""

  def __init__(
      self,
      tool_tag: str,
      invocation_id: str,
      user_name: str,
      source_root: str,
      platform_version: str,
      python_version: str,
      client: clearcut_client.Clearcut,
  ):
    self.tool_tag = tool_tag
    self.invocation_id = invocation_id
    self.user_name = user_name
    self.source_root = source_root
    self.platform_version = platform_version
    self.python_version = python_version
    self._clearcut_client = client

  @classmethod
  def create(cls, tool_tag: str):
    return ToolEventLogger(
        tool_tag=tool_tag,
        invocation_id=str(uuid.uuid4()),
        user_name=getpass.getuser(),
        source_root=os.environ.get('ANDROID_BUILD_TOP', ''),
        platform_version=platform.platform(),
        python_version=platform.python_version(),
        client=clearcut_client.Clearcut(LOG_SOURCE),
    )

  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc_val, exc_tb):
    self.flush()

  def log_invocation_started(self, event_time: datetime, command_args: str):
    """Creates an event log with invocation started info."""
    event = self._create_tool_event()
    event.invocation_started.CopyFrom(
        tool_event_pb2.ToolEvent.InvocationStarted(
            command_args=command_args,
            os=f'{self.platform_version}:{self.python_version}',
        )
    )

    logging.debug('Log invocation_started: %s', event)
    self._log_clearcut_event(event, event_time)

  def log_invocation_stopped(
      self,
      event_time: datetime,
      exit_code: int,
      exit_log: str,
  ):
    """Creates an event log with invocation stopped info."""
    event = self._create_tool_event()
    event.invocation_stopped.CopyFrom(
        tool_event_pb2.ToolEvent.InvocationStopped(
            exit_code=exit_code,
            exit_log=exit_log,
        )
    )

    logging.debug('Log invocation_stopped: %s', event)
    self._log_clearcut_event(event, event_time)

  def flush(self):
    """Sends all batched events to Clearcut."""
    logging.debug('Sending events to Clearcut.')
    self._clearcut_client.flush_events()

  def _create_tool_event(self):
    return tool_event_pb2.ToolEvent(
        tool_tag=self.tool_tag,
        invocation_id=self.invocation_id,
        user_name=self.user_name,
        source_root=self.source_root,
    )

  def _log_clearcut_event(
      self, tool_event: tool_event_pb2.ToolEvent, event_time: datetime
  ):
    log_event = clientanalytics_pb2.LogEvent(
        event_time_ms=int(event_time.timestamp() * 1000),
        source_extension=tool_event.SerializeToString(),
    )
    self._clearcut_client.log(log_event)


class ArgumentParserWithLogging(argparse.ArgumentParser):

  def error(self, message):
    logging.error('Failed to parse args with error: %s', message)
    super().error(message)


def create_arg_parser():
  """Creates an instance of the default ToolEventLogger arg parser."""

  parser = ArgumentParserWithLogging(
      description='Build and upload logs for Android dev tools',
      add_help=True,
      formatter_class=argparse.RawDescriptionHelpFormatter,
  )

  parser.add_argument(
      '--tool_tag',
      type=str,
      required=True,
      help='Name of the tool.',
  )

  parser.add_argument(
      '--start_timestamp',
      type=lambda ts: datetime.datetime.fromtimestamp(float(ts)),
      required=True,
      help=(
          'Timestamp when the tool starts. The timestamp should have the format'
          '%s.%N which represents the seconds elapses since epoch.'
      ),
  )

  parser.add_argument(
      '--end_timestamp',
      type=lambda ts: datetime.datetime.fromtimestamp(float(ts)),
      required=True,
      help=(
          'Timestamp when the tool exits. The timestamp should have the format'
          '%s.%N which represents the seconds elapses since epoch.'
      ),
  )

  parser.add_argument(
      '--tool_args',
      type=str,
      help='Parameters that are passed to the tool.',
  )

  parser.add_argument(
      '--exit_code',
      type=int,
      required=True,
      help='Tool exit code.',
  )

  parser.add_argument(
      '--exit_log',
      type=str,
      help='Logs when tool exits.',
  )

  parser.add_argument(
      '--dry_run',
      action='store_true',
      help='Dry run the tool event logger if set.',
  )

  return parser


def configure_logging():
  root_logging_dir = tempfile.mkdtemp(prefix='tool_event_logger_')

  log_fmt = '%(asctime)s %(filename)s:%(lineno)s:%(levelname)s: %(message)s'
  date_fmt = '%Y-%m-%d %H:%M:%S'
  _, log_path = tempfile.mkstemp(dir=root_logging_dir, suffix='.log')

  logging.basicConfig(
      filename=log_path, level=logging.DEBUG, format=log_fmt, datefmt=date_fmt
  )


def main(argv: list[str]):
  args = create_arg_parser().parse_args(argv[1:])

  if args.dry_run:
    logging.debug('This is a dry run.')
    return

  try:
    with ToolEventLogger.create(args.tool_tag) as logger:
      logger.log_invocation_started(args.start_timestamp, args.tool_args)
      logger.log_invocation_stopped(
          args.end_timestamp, args.exit_code, args.exit_log
      )
  except Exception as e:
    logging.error('Log failed with unexpected error: %s', e)
    raise


if __name__ == '__main__':
  configure_logging()
  main(sys.argv)
+209 −0
Original line number Diff line number Diff line
# Copyright 2024, The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unittests for ToolEventLogger."""

import datetime
import logging
import unittest
from unittest import mock

from atest.metrics import clearcut_client
from proto import tool_event_pb2
from tool_event_logger import tool_event_logger

TEST_INVOCATION_ID = 'test_invocation_id'
TEST_USER_NAME = 'test_user'
TEST_TOOL_TAG = 'test_tool'
TEST_SOURCE_ROOT = 'test_source_root'
TEST_PLATFORM_VERSION = 'test_platform_version'
TEST_PYTHON_VERSION = 'test_python_version'
TEST_EVENT_TIMESTAMP = datetime.datetime.now()


class ToolEventLoggerTest(unittest.TestCase):

  def setUp(self):
    super().setUp()
    self.clearcut_client = FakeClearcutClient()
    self.logger = tool_event_logger.ToolEventLogger(
        TEST_TOOL_TAG,
        TEST_INVOCATION_ID,
        TEST_USER_NAME,
        TEST_SOURCE_ROOT,
        TEST_PLATFORM_VERSION,
        TEST_PYTHON_VERSION,
        client=self.clearcut_client,
    )

  def test_log_event_timestamp(self):
    with self.logger:
      self.logger.log_invocation_started(
          datetime.datetime.fromtimestamp(100.101), 'test_command'
      )

    self.assertEqual(
        self.clearcut_client.get_last_sent_event().event_time_ms, 100101
    )

  def test_log_event_basic_information(self):
    with self.logger:
      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')

    sent_event = self.clearcut_client.get_last_sent_event()
    log_event = tool_event_pb2.ToolEvent.FromString(sent_event.source_extension)
    self.assertEqual(log_event.invocation_id, TEST_INVOCATION_ID)
    self.assertEqual(log_event.user_name, TEST_USER_NAME)
    self.assertEqual(log_event.tool_tag, TEST_TOOL_TAG)
    self.assertEqual(log_event.source_root, TEST_SOURCE_ROOT)

  def test_log_invocation_started(self):
    expected_invocation_started = tool_event_pb2.ToolEvent.InvocationStarted(
        command_args='test_command',
        os=TEST_PLATFORM_VERSION + ':' + TEST_PYTHON_VERSION,
    )

    with self.logger:
      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')

    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
    sent_event = self.clearcut_client.get_last_sent_event()
    self.assertEqual(
        expected_invocation_started,
        tool_event_pb2.ToolEvent.FromString(
            sent_event.source_extension
        ).invocation_started,
    )

  def test_log_invocation_stopped(self):
    expected_invocation_stopped = tool_event_pb2.ToolEvent.InvocationStopped(
        exit_code=0,
        exit_log='exit_log',
    )

    with self.logger:
      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')

    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 1)
    sent_event = self.clearcut_client.get_last_sent_event()
    self.assertEqual(
        expected_invocation_stopped,
        tool_event_pb2.ToolEvent.FromString(
            sent_event.source_extension
        ).invocation_stopped,
    )

  def test_log_multiple_events(self):
    with self.logger:
      self.logger.log_invocation_started(TEST_EVENT_TIMESTAMP, 'test_command')
      self.logger.log_invocation_stopped(TEST_EVENT_TIMESTAMP, 0, 'exit_log')

    self.assertEqual(self.clearcut_client.get_number_of_sent_events(), 2)


class MainTest(unittest.TestCase):

  REQUIRED_ARGS = [
      '',
      '--tool_tag',
      'test_tool',
      '--start_timestamp',
      '1',
      '--end_timestamp',
      '2',
      '--exit_code',
      '0',
  ]

  def test_log_and_exit_with_missing_required_args(self):
    with self.assertLogs() as logs:
      with self.assertRaises(SystemExit) as ex:
        tool_event_logger.main(['', '--tool_tag', 'test_tool'])

    with self.subTest('Verify exception code'):
      self.assertEqual(ex.exception.code, 2)

    with self.subTest('Verify log messages'):
      self.assertIn(
          'the following arguments are required',
          '\n'.join(logs.output),
      )

  def test_log_and_exit_with_invalid_args(self):
    with self.assertLogs() as logs:
      with self.assertRaises(SystemExit) as ex:
        tool_event_logger.main(['', '--start_timestamp', 'test'])

    with self.subTest('Verify exception code'):
      self.assertEqual(ex.exception.code, 2)

    with self.subTest('Verify log messages'):
      self.assertIn(
          '--start_timestamp: invalid',
          '\n'.join(logs.output),
      )

  def test_log_and_exit_with_dry_run(self):
    with self.assertLogs(level=logging.DEBUG) as logs:
      tool_event_logger.main(self.REQUIRED_ARGS + ['--dry_run'])

    with self.subTest('Verify log messages'):
      self.assertIn('dry run', '\n'.join(logs.output))

  @mock.patch.object(clearcut_client, 'Clearcut')
  def test_log_and_exit_with_unexpected_exception(self, mock_cc):
    mock_cc.return_value = FakeClearcutClient(raise_log_exception=True)

    with self.assertLogs() as logs:
      with self.assertRaises(Exception) as ex:
        tool_event_logger.main(self.REQUIRED_ARGS)

    with self.subTest('Verify log messages'):
      self.assertIn('unexpected error', '\n'.join(logs.output))

  @mock.patch.object(clearcut_client, 'Clearcut')
  def test_success(self, mock_cc):
    mock_clear_cut_client = FakeClearcutClient()
    mock_cc.return_value = mock_clear_cut_client

    tool_event_logger.main(self.REQUIRED_ARGS)

    self.assertEqual(mock_clear_cut_client.get_number_of_sent_events(), 2)


class FakeClearcutClient:

  def __init__(self, raise_log_exception=False):
    self.pending_log_events = []
    self.sent_log_events = []
    self.raise_log_exception = raise_log_exception

  def log(self, log_event):
    if self.raise_log_exception:
      raise Exception('unknown exception')
    self.pending_log_events.append(log_event)

  def flush_events(self):
    self.sent_log_events.extend(self.pending_log_events)
    self.pending_log_events.clear()

  def get_number_of_sent_events(self):
    return len(self.sent_log_events)

  def get_last_sent_event(self):
    return self.sent_log_events[-1]


if __name__ == '__main__':
  unittest.main()