build/util/lib/proto/time_consumption_unittests.py

#!/usr/bin/env vpython3

# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""File for testing time_consumption.py."""

import os
import tempfile
import unittest
import unittest.mock as mock

import google.protobuf.json_format as json_format

from google.protobuf import any_pb2

import measures

from test_script_metrics_pb2 import TestScriptMetric, TestScriptMetrics


class TimeConsumptionTest(unittest.TestCase):
  """Test time_consumption.py."""

  @mock.patch('time_consumption.time.time', side_effect=[100, 110])
  def test_by_dumping(self, time_patch) -> None:
    before = len(measures._metric._metrics)
    with measures.time_consumption('test', 'time', 'consumption'):
      pass
    with tempfile.TemporaryDirectory() as tmpdir:
      measures.dump(tmpdir)
      with open(os.path.join(tmpdir,
                             measures.TEST_SCRIPT_METRICS_JSONPB_FILENAME),
                'r',
                encoding='utf-8') as rf:
        any_msg = json_format.Parse(rf.read(), any_pb2.Any())
        message = TestScriptMetrics()
        self.assertTrue(any_msg.Unpack(message))
    self.assertEqual(len(message.metrics), before + 1)
    exp = TestScriptMetric()
    exp.name = 'test/time/consumption (seconds)'
    exp.value = 10
    self.assertEqual(message.metrics[-1], exp)
    self.assertEqual(time_patch.call_count, 2)

  @mock.patch('time_consumption.time.time', side_effect=[100, 101, 102, 110])
  def test_exit_twice(self, time_patch) -> None:
    # This is not a common use scenario, but it shouldn't crash.
    before = len(measures._metric._metrics)
    consumption = measures.time_consumption('test', 'time', 'consumption2')
    with consumption:
      pass
    with consumption:
      pass
    with tempfile.TemporaryDirectory() as tmpdir:
      measures.dump(tmpdir)
      with open(os.path.join(tmpdir,
                             measures.TEST_SCRIPT_METRICS_JSONPB_FILENAME),
                'r',
                encoding='utf-8') as rf:
        any_msg = json_format.Parse(rf.read(), any_pb2.Any())
        message = TestScriptMetrics()
        self.assertTrue(any_msg.Unpack(message))
    self.assertEqual(len(message.metrics), before + 1)
    exp = TestScriptMetric()
    exp.name = 'test/time/consumption2 (seconds)'
    exp.value = 8
    self.assertEqual(message.metrics[-1], exp)
    self.assertEqual(time_patch.call_count, 4)


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