third_party/pigweed/repo/pw_protobuf_compiler/py/python_protos_test.py

   1 #!/usr/bin/env python3
   2 # Copyright 2020 The Pigweed Authors
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   5 # use this file except in compliance with the License. You may obtain a copy of
   6 # the License at
   7 #
   8 #     https://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  12 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  13 # License for the specific language governing permissions and limitations under
  14 # the License.
  15 """Tests compiling and importing Python protos on the fly."""
  16
  17 from pathlib import Path
  18 import tempfile
  19 import unittest
  20
  21 from pw_protobuf_compiler import python_protos
  22
  23 PROTO_1 = """\
  24 syntax = "proto3";
  25
  26 package pw.protobuf_compiler.test1;
  27
  28 message SomeMessage {
  29   uint32 magic_number = 1;
  30 }
  31
  32 message AnotherMessage {
  33   enum Result {
  34     FAILED = 0;
  35     FAILED_MISERABLY = 1;
  36     I_DONT_WANT_TO_TALK_ABOUT_IT = 2;
  37   }
  38
  39   Result result = 1;
  40   string payload = 2;
  41 }
  42
  43 service PublicService {
  44   rpc Unary(SomeMessage) returns (AnotherMessage) {}
  45   rpc ServerStreaming(SomeMessage) returns (stream AnotherMessage) {}
  46   rpc ClientStreaming(stream SomeMessage) returns (AnotherMessage) {}
  47   rpc BidiStreaming(stream SomeMessage) returns (stream AnotherMessage) {}
  48 }
  49 """
  50
  51 PROTO_2 = """\
  52 syntax = "proto2";
  53
  54 package pw.protobuf_compiler.test2;
  55
  56 message Request {
  57   optional float magic_number = 1;
  58 }
  59
  60 message Response {
  61 }
  62
  63 service Alpha {
  64   rpc Unary(Request) returns (Response) {}
  65 }
  66
  67 service Bravo {
  68   rpc BidiStreaming(stream Request) returns (stream Response) {}
  69 }
  70 """
  71
  72 PROTO_3 = """\
  73 syntax = "proto3";
  74
  75 package pw.protobuf_compiler.test2;
  76
  77 enum Greeting {
  78   YO = 0;
  79   HI = 1;
  80 }
  81
  82 message Hello {
  83   repeated int64 value = 1;
  84   Greeting hi = 2;
  85 }
  86 """
  87
  88
  89 class TestCompileAndImport(unittest.TestCase):
  90     def setUp(self):
  91         self._proto_dir = tempfile.TemporaryDirectory(prefix='proto_test')
  92         self._protos = []
  93
  94         for i, contents in enumerate([PROTO_1, PROTO_2, PROTO_3], 1):
  95             self._protos.append(Path(self._proto_dir.name, f'test_{i}.proto'))
  96             self._protos[-1].write_text(contents)
  97
  98     def tearDown(self):
  99         self._proto_dir.cleanup()
 100
 101     def test_compile_to_temp_dir_and_import(self):
 102         modules = {
 103             m.DESCRIPTOR.name: m
 104             for m in python_protos.compile_and_import(self._protos)
 105         }
 106         self.assertEqual(3, len(modules))
 107
 108         # Make sure the protobuf modules contain what we expect them to.
 109         mod = modules['test_1.proto']
 110         self.assertEqual(
 111             4, len(mod.DESCRIPTOR.services_by_name['PublicService'].methods))
 112
 113         mod = modules['test_2.proto']
 114         self.assertEqual(mod.Request(magic_number=1.5).magic_number, 1.5)
 115         self.assertEqual(2, len(mod.DESCRIPTOR.services_by_name))
 116
 117         mod = modules['test_3.proto']
 118         self.assertEqual(mod.Hello(value=[123, 456]).value, [123, 456])
 119
 120
 121 class TestProtoLibrary(TestCompileAndImport):
 122     """Tests the Library class."""
 123     def setUp(self):
 124         super().setUp()
 125         self._library = python_protos.Library(
 126             python_protos.compile_and_import(self._protos))
 127
 128     def test_packages_can_access_messages(self):
 129         msg = self._library.packages.pw.protobuf_compiler.test1.SomeMessage
 130         self.assertEqual(msg(magic_number=123).magic_number, 123)
 131
 132     def test_packages_finds_across_modules(self):
 133         msg = self._library.packages.pw.protobuf_compiler.test2.Request
 134         self.assertEqual(msg(magic_number=50).magic_number, 50)
 135
 136         val = self._library.packages.pw.protobuf_compiler.test2.YO
 137         self.assertEqual(val, 0)
 138
 139     def test_packages_invalid_name(self):
 140         with self.assertRaises(AttributeError):
 141             _ = self._library.packages.nothing
 142
 143         with self.assertRaises(AttributeError):
 144             _ = self._library.packages.pw.NOT_HERE
 145
 146         with self.assertRaises(AttributeError):
 147             _ = self._library.packages.pw.protobuf_compiler.test1.NotARealMsg
 148
 149     def test_access_modules_by_package(self):
 150         test1 = self._library.modules_by_package['pw.protobuf_compiler.test1']
 151         self.assertEqual(len(test1), 1)
 152         self.assertEqual(test1[0].AnotherMessage.Result.Value('FAILED'), 0)
 153
 154         test2 = self._library.modules_by_package['pw.protobuf_compiler.test2']
 155         self.assertEqual(len(test2), 2)
 156
 157     def test_access_modules_by_package_unknown(self):
 158         with self.assertRaises(KeyError):
 159             _ = self._library.modules_by_package['pw.not_real']
 160
 161     def test_library_from_strings(self):
 162         # Replace the package to avoid conflicts with the other proto imports
 163         new_protos = [
 164             p.replace('pw.protobuf_compiler', 'proto.library.test')
 165             for p in [PROTO_1, PROTO_2, PROTO_3]
 166         ]
 167
 168         library = python_protos.Library.from_strings(new_protos)
 169
 170         # Make sure we can safely import the same proto contents multiple times.
 171         library = python_protos.Library.from_strings(new_protos)
 172
 173         msg = library.packages.proto.library.test.test2.Request
 174         self.assertEqual(msg(magic_number=50).magic_number, 50)
 175
 176         val = library.packages.proto.library.test.test2.YO
 177         self.assertEqual(val, 0)
 178
 179     def test_access_nested_packages_by_name(self):
 180         self.assertIs(self._library.packages['pw.protobuf_compiler.test1'],
 181                       self._library.packages.pw.protobuf_compiler.test1)
 182         self.assertIs(self._library.packages.pw['protobuf_compiler.test1'],
 183                       self._library.packages.pw.protobuf_compiler.test1)
 184         self.assertIs(self._library.packages.pw.protobuf_compiler['test1'],
 185                       self._library.packages.pw.protobuf_compiler.test1)
 186
 187     def test_access_nested_packages_by_name_unknown_package(self):
 188         with self.assertRaises(KeyError):
 189             _ = self._library.packages['']
 190
 191         with self.assertRaises(KeyError):
 192             _ = self._library.packages['.']
 193
 194         with self.assertRaises(KeyError):
 195             _ = self._library.packages['protobuf_compiler.test1']
 196
 197         with self.assertRaises(KeyError):
 198             _ = self._library.packages.pw['pw.protobuf_compiler.test1']
 199
 200         with self.assertRaises(KeyError):
 201             _ = self._library.packages.pw.protobuf_compiler['not here']
 202
 203
 204 if __name__ == '__main__':
 205     unittest.main()