blob: 139800d99147c198fef8a0aedc0ba013935b91a9 [file] [log] [blame]
#!/usr/bin/env python3
# Copyright 2020 The Pigweed Authors
#
# 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
#
# https://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.
"""Tests compiling and importing Python protos on the fly."""
from pathlib import Path
import tempfile
import unittest
from pw_protobuf_compiler import python_protos
from pw_protobuf_compiler.python_protos import bytes_repr, proto_repr
PROTO_1 = """\
syntax = "proto3";
package pw.protobuf_compiler.test1;
message SomeMessage {
uint32 magic_number = 1;
}
message AnotherMessage {
enum Result {
FAILED = 0;
FAILED_MISERABLY = 1;
I_DONT_WANT_TO_TALK_ABOUT_IT = 2;
}
Result result = 1;
string payload = 2;
}
service PublicService {
rpc Unary(SomeMessage) returns (AnotherMessage) {}
rpc ServerStreaming(SomeMessage) returns (stream AnotherMessage) {}
rpc ClientStreaming(stream SomeMessage) returns (AnotherMessage) {}
rpc BidiStreaming(stream SomeMessage) returns (stream AnotherMessage) {}
}
"""
PROTO_2 = """\
syntax = "proto2";
package pw.protobuf_compiler.test2;
message Request {
optional float magic_number = 1;
}
message Response {
}
service Alpha {
rpc Unary(Request) returns (Response) {}
}
service Bravo {
rpc BidiStreaming(stream Request) returns (stream Response) {}
}
"""
PROTO_3 = """\
syntax = "proto3";
package pw.protobuf_compiler.test2;
enum Greeting {
YO = 0;
HI = 1;
}
message Hello {
repeated int64 value = 1;
Greeting hi = 2;
}
message NestingMessage {
message NestedMessage {
message NestedNestedMessage {
int32 nested_nested_field = 1;
}
NestedNestedMessage nested_nested_message = 1;
}
NestedMessage nested_message = 1;
}
"""
class TestCompileAndImport(unittest.TestCase):
"""Test compiling and importing."""
def setUp(self):
self._proto_dir = tempfile.TemporaryDirectory(prefix='proto_test')
self._protos = []
for i, contents in enumerate([PROTO_1, PROTO_2, PROTO_3], 1):
self._protos.append(Path(self._proto_dir.name, f'test_{i}.proto'))
self._protos[-1].write_text(contents)
def tearDown(self):
self._proto_dir.cleanup()
def test_compile_to_temp_dir_and_import(self):
modules = {
m.DESCRIPTOR.name: m
for m in python_protos.compile_and_import(self._protos)
}
self.assertEqual(3, len(modules))
# Make sure the protobuf modules contain what we expect them to.
mod = modules['test_1.proto']
self.assertEqual(
4, len(mod.DESCRIPTOR.services_by_name['PublicService'].methods))
mod = modules['test_2.proto']
self.assertEqual(mod.Request(magic_number=1.5).magic_number, 1.5)
self.assertEqual(2, len(mod.DESCRIPTOR.services_by_name))
mod = modules['test_3.proto']
self.assertEqual(mod.Hello(value=[123, 456]).value, [123, 456])
class TestProtoLibrary(TestCompileAndImport):
"""Tests the Library class."""
def setUp(self):
super().setUp()
self._library = python_protos.Library(
python_protos.compile_and_import(self._protos))
def test_packages_can_access_messages(self):
msg = self._library.packages.pw.protobuf_compiler.test1.SomeMessage
self.assertEqual(msg(magic_number=123).magic_number, 123)
def test_packages_finds_across_modules(self):
msg = self._library.packages.pw.protobuf_compiler.test2.Request
self.assertEqual(msg(magic_number=50).magic_number, 50)
val = self._library.packages.pw.protobuf_compiler.test2.YO
self.assertEqual(val, 0)
def test_packages_invalid_name(self):
with self.assertRaises(AttributeError):
_ = self._library.packages.nothing
with self.assertRaises(AttributeError):
_ = self._library.packages.pw.NOT_HERE
with self.assertRaises(AttributeError):
_ = self._library.packages.pw.protobuf_compiler.test1.NotARealMsg
def test_access_modules_by_package(self):
test1 = self._library.modules_by_package['pw.protobuf_compiler.test1']
self.assertEqual(len(test1), 1)
self.assertEqual(test1[0].AnotherMessage.Result.Value('FAILED'), 0)
test2 = self._library.modules_by_package['pw.protobuf_compiler.test2']
self.assertEqual(len(test2), 2)
def test_access_modules_by_package_unknown(self):
with self.assertRaises(KeyError):
_ = self._library.modules_by_package['pw.not_real']
def test_library_from_strings(self):
# Replace the package to avoid conflicts with the other proto imports
new_protos = [
p.replace('pw.protobuf_compiler', 'proto.library.test')
for p in [PROTO_1, PROTO_2, PROTO_3]
]
library = python_protos.Library.from_strings(new_protos)
# Make sure we can safely import the same proto contents multiple times.
library = python_protos.Library.from_strings(new_protos)
msg = library.packages.proto.library.test.test2.Request
self.assertEqual(msg(magic_number=50).magic_number, 50)
val = library.packages.proto.library.test.test2.YO
self.assertEqual(val, 0)
def test_access_nested_packages_by_name(self):
self.assertIs(self._library.packages['pw.protobuf_compiler.test1'],
self._library.packages.pw.protobuf_compiler.test1)
self.assertIs(self._library.packages.pw['protobuf_compiler.test1'],
self._library.packages.pw.protobuf_compiler.test1)
self.assertIs(self._library.packages.pw.protobuf_compiler['test1'],
self._library.packages.pw.protobuf_compiler.test1)
def test_access_nested_packages_by_name_unknown_package(self):
with self.assertRaises(KeyError):
_ = self._library.packages['']
with self.assertRaises(KeyError):
_ = self._library.packages['.']
with self.assertRaises(KeyError):
_ = self._library.packages['protobuf_compiler.test1']
with self.assertRaises(KeyError):
_ = self._library.packages.pw['pw.protobuf_compiler.test1']
with self.assertRaises(KeyError):
_ = self._library.packages.pw.protobuf_compiler['not here']
def test_messages(self):
protos = self._library.packages.pw.protobuf_compiler
self.assertEqual(
set(self._library.messages()), {
protos.test1.SomeMessage,
protos.test1.AnotherMessage,
protos.test2.Request,
protos.test2.Response,
protos.test2.Hello,
protos.test2.NestingMessage,
protos.test2.NestingMessage.NestedMessage,
protos.test2.NestingMessage.NestedMessage.NestedNestedMessage,
})
PROTO_FOR_REPR = """\
syntax = "proto3";
package pw.test3;
enum Enum {
ZERO = 0;
ONE = 1;
}
message Nested {
repeated int64 value = 1;
Enum an_enum = 2;
}
message Message {
Nested message = 1;
repeated Nested repeated_message = 2;
fixed32 regular_int = 3;
optional int64 optional_int = 4;
repeated int32 repeated_int = 5;
bytes regular_bytes = 6;
optional bytes optional_bytes = 7;
repeated bytes repeated_bytes = 8;
string regular_string = 9;
optional string optional_string = 10;
repeated string repeated_string = 11;
Enum regular_enum = 12;
optional Enum optional_enum = 13;
repeated Enum repeated_enum = 14;
oneof oneof_test {
string oneof_1 = 15;
int32 oneof_2 = 16;
Nested oneof_3 = 17;
}
map<string, Nested> mapping = 18;
}
"""
class TestProtoRepr(unittest.TestCase):
"""Tests printing protobufs."""
def setUp(self):
protos = python_protos.Library.from_strings(PROTO_FOR_REPR)
self.enum = protos.packages.pw.test3.Enum
self.nested = protos.packages.pw.test3.Nested
self.message = protos.packages.pw.test3.Message
def test_empty(self):
self.assertEqual('pw.test3.Nested()', proto_repr(self.nested()))
self.assertEqual('pw.test3.Message()', proto_repr(self.message()))
def test_int_fields(self):
self.assertEqual(
'pw.test3.Message('
'regular_int=999, '
'optional_int=-1, '
'repeated_int=[0, 1, 2])',
proto_repr(self.message(repeated_int=[0, 1, 2],
regular_int=999,
optional_int=-1),
wrap=False))
def test_bytes_fields(self):
self.assertEqual(
'pw.test3.Message('
r"regular_bytes=b'\xFE\xED\xBE\xEF', "
r"optional_bytes=b'', "
r"repeated_bytes=[b'Hello\'\'\''])",
proto_repr(self.message(
regular_bytes=b'\xfe\xed\xbe\xef',
optional_bytes=b'',
repeated_bytes=[b"Hello'''"],
),
wrap=False))
def test_string_fields(self):
self.assertEqual(
'pw.test3.Message('
"regular_string='hi', "
"optional_string='', "
'repeated_string=["\'"])',
proto_repr(self.message(
regular_string='hi',
optional_string='',
repeated_string=[b"'"],
),
wrap=False))
def test_enum_fields(self):
self.assertEqual('pw.test3.Nested(an_enum=pw.test3.Enum.ONE)',
proto_repr(self.nested(an_enum=1)))
self.assertEqual('pw.test3.Message(optional_enum=pw.test3.Enum.ONE)',
proto_repr(self.message(optional_enum=self.enum.ONE)))
self.assertEqual(
'pw.test3.Message(repeated_enum='
'[pw.test3.Enum.ONE, pw.test3.Enum.ONE, pw.test3.Enum.ZERO])',
proto_repr(self.message(repeated_enum=[1, 1, 0]), wrap=False))
def test_message_fields(self):
self.assertEqual(
'pw.test3.Message(message=pw.test3.Nested(value=[123]))',
proto_repr(self.message(message=self.nested(value=[123]))))
self.assertEqual(
'pw.test3.Message('
'repeated_message=[pw.test3.Nested(value=[123]), '
'pw.test3.Nested()])',
proto_repr(self.message(
repeated_message=[self.nested(
value=[123]), self.nested()]),
wrap=False))
def test_optional_shown_if_set_to_default(self):
self.assertEqual(
"pw.test3.Message("
"optional_int=0, optional_bytes=b'', optional_string='', "
"optional_enum=pw.test3.Enum.ZERO)",
proto_repr(self.message(optional_int=0,
optional_bytes=b'',
optional_string='',
optional_enum=0),
wrap=False))
def test_oneof(self):
self.assertEqual(proto_repr(self.message(oneof_1='test')),
"pw.test3.Message(oneof_1='test')")
self.assertEqual(proto_repr(self.message(oneof_2=123)),
"pw.test3.Message(oneof_2=123)")
self.assertEqual(
proto_repr(
self.message(oneof_3=self.nested(an_enum=self.enum.ONE))),
'pw.test3.Message('
'oneof_3=pw.test3.Nested(an_enum=pw.test3.Enum.ONE))')
msg = self.message(oneof_1='test')
msg.oneof_2 = 99
self.assertEqual(proto_repr(msg), "pw.test3.Message(oneof_2=99)")
def test_map(self):
msg = self.message()
msg.mapping['zero'].MergeFrom(self.nested())
msg.mapping['one'].MergeFrom(
self.nested(an_enum=self.enum.ONE, value=[1]))
result = proto_repr(msg, wrap=False)
self.assertRegex(result, r'^pw.test3.Message\(mapping={.*}\)$')
self.assertIn("'zero': pw.test3.Nested()", result)
self.assertIn(
"'one': pw.test3.Nested(value=[1], an_enum=pw.test3.Enum.ONE)",
result)
def test_bytes_repr(self):
self.assertEqual(bytes_repr(b'\xfe\xed\xbe\xef'),
r"b'\xFE\xED\xBE\xEF'")
self.assertEqual(bytes_repr(b'\xfe\xed\xbe\xef123'),
r"b'\xFE\xED\xBE\xEF\x31\x32\x33'")
self.assertEqual(bytes_repr(b'\xfe\xed\xbe\xef1234'),
r"b'\xFE\xED\xBE\xEF1234'")
self.assertEqual(bytes_repr(b'\xfe\xed\xbe\xef12345'),
r"b'\xFE\xED\xBE\xEF12345'")
def test_wrap_multiple_lines(self):
self.assertEqual(
"""\
pw.test3.Message(optional_int=0,
optional_bytes=b'',
optional_string='',
optional_enum=pw.test3.Enum.ZERO)""",
proto_repr(self.message(optional_int=0,
optional_bytes=b'',
optional_string='',
optional_enum=0),
wrap=True))
def test_wrap_one_line(self):
self.assertEqual(
"pw.test3.Message(optional_int=0, optional_bytes=b'')",
proto_repr(self.message(optional_int=0, optional_bytes=b''),
wrap=True))
if __name__ == '__main__':
unittest.main()