build/rbe/pb_message_util_test.py

#!/usr/bin/env fuchsia-vendored-python
# Copyright 2022 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import unittest
from unittest import mock

import pb_message_util

# The function being tested works on any subclass of message.Message
# which is an abstract base class.  However, we need a concrete message type
# to test.  We choose api.log.log_pb2.LogDump because this is what
# upload_reproxy_logs.py uses, and it is very stable.
from api.log import log_pb2
from go.api.command import command_pb2
from google.protobuf import timestamp_pb2


class ProtoMessageToBQDictTest(unittest.TestCase):
    def test_empty_log(self):
        log = log_pb2.LogDump()
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(str(converted_log), "{}")

    def test_one_record(self):
        log = log_pb2.LogDump(records=[log_pb2.LogRecord()])
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(str(converted_log), "{'records': [{}]}")

    def test_two_records(self):
        log = log_pb2.LogDump(
            records=[log_pb2.LogRecord(), log_pb2.LogRecord()]
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(str(converted_log), "{'records': [{}, {}]}")

    def test_nested_message_strings(self):
        log = log_pb2.LogDump(
            records=[
                log_pb2.LogRecord(
                    command=command_pb2.Command(
                        exec_root="/home",
                        args=["echo", "hello"],
                    )
                )
            ]
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(
            str(converted_log),
            "{'records': [{'command': {'exec_root': '/home', 'args': ['echo', 'hello']}}]}",
        )

    def test_nested_message_map_string_string(self):
        """Exercises protobuf.containers.ScalarMap"""
        log = log_pb2.LogDump(
            records=[
                log_pb2.LogRecord(
                    command=command_pb2.Command(
                        platform={"foo": "bar", "baz": "quux"},
                    )
                )
            ]
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(
            str(converted_log),
            "{'records': [{'command': {'platform': [{'key': 'baz', 'value': 'quux'}, {'key': 'foo', 'value': 'bar'}]}}]}",
        )

    def test_nested_message_map_string_message(self):
        """Exercises protobuf.containers.MessageMap"""
        log = log_pb2.LogDump(
            records=[
                log_pb2.LogRecord(
                    remote_metadata=log_pb2.RemoteMetadata(
                        event_times={
                            "birth": command_pb2.TimeInterval(
                                to=timestamp_pb2.Timestamp(seconds=1601428537)
                            ),
                            "death": command_pb2.TimeInterval(
                                to=timestamp_pb2.Timestamp(seconds=1601439840)
                            ),
                        },
                    )
                )
            ]
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(
            str(converted_log),
            "{'records': [{'remote_metadata': {'event_times': [{'key': 'birth', 'value': {'to': '2020-09-30T01:15:37Z'}}, {'key': 'death', 'value': {'to': '2020-09-30T04:24:00Z'}}]}}]}",
        )

    def test_nested_message_enum(self):
        """Make sure enum name (not number) is printed."""
        log = log_pb2.LogDump(
            records=[
                log_pb2.LogRecord(
                    completion_status=log_pb2.CompletionStatus.STATUS_CACHE_HIT
                )
            ]
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(
            str(converted_log),
            "{'records': [{'completion_status': 'STATUS_CACHE_HIT'}]}",
        )

    def test_timestamp(self):
        """Make sure timestamp is formatted, and not (seconds, nanos)"""
        log = command_pb2.TimeInterval(
            to=timestamp_pb2.Timestamp(seconds=1661472513, nanos=114242156)
        )
        converted_log = pb_message_util.proto_message_to_bq_dict(log)
        self.assertEqual(
            str(converted_log), "{'to': '2022-08-26T00:08:33.114242156Z'}"
        )


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