1#
2# Copyright 2015 Google Inc.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://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,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import base64
17import datetime
18import json
19import sys
20
21import unittest2
22
23from apitools.base.protorpclite import message_types
24from apitools.base.protorpclite import messages
25from apitools.base.protorpclite import util
26from apitools.base.py import encoding
27from apitools.base.py import exceptions
28from apitools.base.py import extra_types
29
30
31class SimpleMessage(messages.Message):
32    field = messages.StringField(1)
33    repfield = messages.StringField(2, repeated=True)
34
35
36class BytesMessage(messages.Message):
37    field = messages.BytesField(1)
38    repfield = messages.BytesField(2, repeated=True)
39
40
41class TimeMessage(messages.Message):
42    timefield = message_types.DateTimeField(3)
43
44
45@encoding.MapUnrecognizedFields('additional_properties')
46class AdditionalPropertiesMessage(messages.Message):
47
48    class AdditionalProperty(messages.Message):
49        key = messages.StringField(1)
50        value = messages.StringField(2)
51
52    additional_properties = messages.MessageField(
53        'AdditionalProperty', 1, repeated=True)
54
55
56@encoding.MapUnrecognizedFields('additional_properties')
57class AdditionalIntPropertiesMessage(messages.Message):
58
59    class AdditionalProperty(messages.Message):
60        key = messages.StringField(1)
61        value = messages.IntegerField(2)
62
63    additional_properties = messages.MessageField(
64        'AdditionalProperty', 1, repeated=True)
65
66
67@encoding.MapUnrecognizedFields('additional_properties')
68class UnrecognizedEnumMessage(messages.Message):
69
70    class ThisEnum(messages.Enum):
71        VALUE_ONE = 1
72        VALUE_TWO = 2
73
74    class AdditionalProperty(messages.Message):
75        key = messages.StringField(1)
76        value = messages.EnumField('UnrecognizedEnumMessage.ThisEnum', 2)
77
78    additional_properties = messages.MessageField(
79        AdditionalProperty, 1, repeated=True)
80
81
82class CompoundPropertyType(messages.Message):
83    index = messages.IntegerField(1)
84    name = messages.StringField(2)
85
86
87class MessageWithEnum(messages.Message):
88
89    class ThisEnum(messages.Enum):
90        VALUE_ONE = 1
91        VALUE_TWO = 2
92
93    field_one = messages.EnumField(ThisEnum, 1)
94    field_two = messages.EnumField(ThisEnum, 2, default=ThisEnum.VALUE_TWO)
95    ignored_field = messages.EnumField(ThisEnum, 3)
96
97
98@encoding.MapUnrecognizedFields('additional_properties')
99class AdditionalMessagePropertiesMessage(messages.Message):
100
101    class AdditionalProperty(messages.Message):
102        key = messages.StringField(1)
103        value = messages.MessageField(CompoundPropertyType, 2)
104
105    additional_properties = messages.MessageField(
106        'AdditionalProperty', 1, repeated=True)
107
108
109class HasNestedMessage(messages.Message):
110    nested = messages.MessageField(AdditionalPropertiesMessage, 1)
111    nested_list = messages.StringField(2, repeated=True)
112
113
114class ExtraNestedMessage(messages.Message):
115    nested = messages.MessageField(HasNestedMessage, 1)
116
117
118class MessageWithRemappings(messages.Message):
119
120    class SomeEnum(messages.Enum):
121        enum_value = 1
122        second_value = 2
123
124    enum_field = messages.EnumField(SomeEnum, 1)
125    double_encoding = messages.EnumField(SomeEnum, 2)
126    another_field = messages.StringField(3)
127    repeated_enum = messages.EnumField(SomeEnum, 4, repeated=True)
128    repeated_field = messages.StringField(5, repeated=True)
129
130
131@encoding.MapUnrecognizedFields('additional_properties')
132class RepeatedJsonValueMessage(messages.Message):
133
134    class AdditionalProperty(messages.Message):
135        key = messages.StringField(1)
136        value = messages.MessageField(extra_types.JsonValue, 2, repeated=True)
137
138    additional_properties = messages.MessageField('AdditionalProperty', 1,
139                                                  repeated=True)
140
141
142encoding.AddCustomJsonEnumMapping(MessageWithRemappings.SomeEnum,
143                                  'enum_value', 'wire_name')
144encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
145                                   'double_encoding', 'doubleEncoding')
146encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
147                                   'another_field', 'anotherField')
148encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
149                                   'repeated_field', 'repeatedField')
150
151
152class EncodingTest(unittest2.TestCase):
153
154    def testCopyProtoMessage(self):
155        msg = SimpleMessage(field='abc')
156        new_msg = encoding.CopyProtoMessage(msg)
157        self.assertEqual(msg.field, new_msg.field)
158        msg.field = 'def'
159        self.assertNotEqual(msg.field, new_msg.field)
160
161    def testBytesEncoding(self):
162        b64_str = 'AAc+'
163        b64_msg = '{"field": "%s"}' % b64_str
164        urlsafe_b64_str = 'AAc-'
165        urlsafe_b64_msg = '{"field": "%s"}' % urlsafe_b64_str
166        data = base64.b64decode(b64_str)
167        msg = BytesMessage(field=data)
168        self.assertEqual(
169            msg, encoding.JsonToMessage(BytesMessage, urlsafe_b64_msg))
170        self.assertEqual(msg, encoding.JsonToMessage(BytesMessage, b64_msg))
171        self.assertEqual(urlsafe_b64_msg, encoding.MessageToJson(msg))
172
173        enc_rep_msg = '{"repfield": ["%(b)s", "%(b)s"]}' % {
174            'b': urlsafe_b64_str}
175        rep_msg = BytesMessage(repfield=[data, data])
176        self.assertEqual(
177            rep_msg, encoding.JsonToMessage(BytesMessage, enc_rep_msg))
178        self.assertEqual(enc_rep_msg, encoding.MessageToJson(rep_msg))
179
180    def testIncludeFields(self):
181        msg = SimpleMessage()
182        self.assertEqual('{}', encoding.MessageToJson(msg))
183        self.assertEqual(
184            '{"field": null}',
185            encoding.MessageToJson(msg, include_fields=['field']))
186        self.assertEqual(
187            '{"repfield": []}',
188            encoding.MessageToJson(msg, include_fields=['repfield']))
189
190    def testNestedIncludeFields(self):
191        msg = HasNestedMessage(
192            nested=AdditionalPropertiesMessage(
193                additional_properties=[]))
194        self.assertEqual(
195            '{"nested": null}',
196            encoding.MessageToJson(msg, include_fields=['nested']))
197        self.assertEqual(
198            '{"nested": {"additional_properties": []}}',
199            encoding.MessageToJson(
200                msg, include_fields=['nested.additional_properties']))
201        msg = ExtraNestedMessage(nested=msg)
202        self.assertEqual(
203            '{"nested": {"nested": null}}',
204            encoding.MessageToJson(msg, include_fields=['nested.nested']))
205        # When clearing 'nested.nested_list', its sibling ('nested.nested')
206        # should remain unaffected.
207        self.assertIn(
208            encoding.MessageToJson(msg, include_fields=['nested.nested_list']),
209            ['{"nested": {"nested": {}, "nested_list": []}}',
210             '{"nested": {"nested_list": [], "nested": {}}}'])
211        self.assertEqual(
212            '{"nested": {"nested": {"additional_properties": []}}}',
213            encoding.MessageToJson(
214                msg, include_fields=['nested.nested.additional_properties']))
215
216    def testAdditionalPropertyMapping(self):
217        msg = AdditionalPropertiesMessage()
218        msg.additional_properties = [
219            AdditionalPropertiesMessage.AdditionalProperty(
220                key='key_one', value='value_one'),
221            AdditionalPropertiesMessage.AdditionalProperty(
222                key='key_two', value='value_two'),
223        ]
224
225        encoded_msg = encoding.MessageToJson(msg)
226        self.assertEqual(
227            {'key_one': 'value_one', 'key_two': 'value_two'},
228            json.loads(encoded_msg))
229
230        new_msg = encoding.JsonToMessage(type(msg), encoded_msg)
231        self.assertEqual(
232            set(('key_one', 'key_two')),
233            set([x.key for x in new_msg.additional_properties]))
234        self.assertIsNot(msg, new_msg)
235
236        new_msg.additional_properties.pop()
237        self.assertEqual(1, len(new_msg.additional_properties))
238        self.assertEqual(2, len(msg.additional_properties))
239
240    def testNumericPropertyName(self):
241        json_msg = '{"nested": {"123": "def"}}'
242        msg = encoding.JsonToMessage(HasNestedMessage, json_msg)
243        self.assertEqual(1, len(msg.nested.additional_properties))
244
245    def testNumericPropertyValue(self):
246        json_msg = '{"key_one": "123"}'
247        msg = encoding.JsonToMessage(AdditionalIntPropertiesMessage, json_msg)
248        self.assertEqual(
249            AdditionalIntPropertiesMessage(
250                additional_properties=[
251                    AdditionalIntPropertiesMessage.AdditionalProperty(
252                        key='key_one', value=123)]),
253            msg)
254
255    def testAdditionalMessageProperties(self):
256        json_msg = '{"input": {"index": 0, "name": "output"}}'
257        result = encoding.JsonToMessage(
258            AdditionalMessagePropertiesMessage, json_msg)
259        self.assertEqual(1, len(result.additional_properties))
260        self.assertEqual(0, result.additional_properties[0].value.index)
261
262    def testUnrecognizedEnum(self):
263        json_msg = '{"input": "VALUE_ONE"}'
264        result = encoding.JsonToMessage(
265            UnrecognizedEnumMessage, json_msg)
266        self.assertEqual(1, len(result.additional_properties))
267        self.assertEqual(UnrecognizedEnumMessage.ThisEnum.VALUE_ONE,
268                         result.additional_properties[0].value)
269
270    def testNestedFieldMapping(self):
271        nested_msg = AdditionalPropertiesMessage()
272        nested_msg.additional_properties = [
273            AdditionalPropertiesMessage.AdditionalProperty(
274                key='key_one', value='value_one'),
275            AdditionalPropertiesMessage.AdditionalProperty(
276                key='key_two', value='value_two'),
277        ]
278        msg = HasNestedMessage(nested=nested_msg)
279
280        encoded_msg = encoding.MessageToJson(msg)
281        self.assertEqual(
282            {'nested': {'key_one': 'value_one', 'key_two': 'value_two'}},
283            json.loads(encoded_msg))
284
285        new_msg = encoding.JsonToMessage(type(msg), encoded_msg)
286        self.assertEqual(
287            set(('key_one', 'key_two')),
288            set([x.key for x in new_msg.nested.additional_properties]))
289
290        new_msg.nested.additional_properties.pop()
291        self.assertEqual(1, len(new_msg.nested.additional_properties))
292        self.assertEqual(2, len(msg.nested.additional_properties))
293
294    def testValidEnums(self):
295        message_json = '{"field_one": "VALUE_ONE"}'
296        message = encoding.JsonToMessage(MessageWithEnum, message_json)
297        self.assertEqual(MessageWithEnum.ThisEnum.VALUE_ONE, message.field_one)
298        self.assertEqual(MessageWithEnum.ThisEnum.VALUE_TWO, message.field_two)
299        self.assertEqual(json.loads(message_json),
300                         json.loads(encoding.MessageToJson(message)))
301
302    def testIgnoredEnums(self):
303        json_with_typo = '{"field_one": "VALUE_OEN"}'
304        message = encoding.JsonToMessage(MessageWithEnum, json_with_typo)
305        self.assertEqual(None, message.field_one)
306        self.assertEqual(('VALUE_OEN', messages.Variant.ENUM),
307                         message.get_unrecognized_field_info('field_one'))
308        self.assertEqual(json.loads(json_with_typo),
309                         json.loads(encoding.MessageToJson(message)))
310
311        empty_json = ''
312        message = encoding.JsonToMessage(MessageWithEnum, empty_json)
313        self.assertEqual(None, message.field_one)
314
315    def testIgnoredEnumsWithDefaults(self):
316        json_with_typo = '{"field_two": "VALUE_OEN"}'
317        message = encoding.JsonToMessage(MessageWithEnum, json_with_typo)
318        self.assertEqual(MessageWithEnum.ThisEnum.VALUE_TWO, message.field_two)
319        self.assertEqual(json.loads(json_with_typo),
320                         json.loads(encoding.MessageToJson(message)))
321
322    def testUnknownNestedRoundtrip(self):
323        json_message = '{"field": "abc", "submessage": {"a": 1, "b": "foo"}}'
324        message = encoding.JsonToMessage(SimpleMessage, json_message)
325        self.assertEqual(json.loads(json_message),
326                         json.loads(encoding.MessageToJson(message)))
327
328    def testJsonDatetime(self):
329        msg = TimeMessage(timefield=datetime.datetime(
330            2014, 7, 2, 23, 33, 25, 541000,
331            tzinfo=util.TimeZoneOffset(datetime.timedelta(0))))
332        self.assertEqual(
333            '{"timefield": "2014-07-02T23:33:25.541000+00:00"}',
334            encoding.MessageToJson(msg))
335
336    def testEnumRemapping(self):
337        msg = MessageWithRemappings(
338            enum_field=MessageWithRemappings.SomeEnum.enum_value)
339        json_message = encoding.MessageToJson(msg)
340        self.assertEqual('{"enum_field": "wire_name"}', json_message)
341        self.assertEqual(
342            msg, encoding.JsonToMessage(MessageWithRemappings, json_message))
343
344    def testRepeatedEnumRemapping(self):
345        msg = MessageWithRemappings(
346            repeated_enum=[
347                MessageWithRemappings.SomeEnum.enum_value,
348                MessageWithRemappings.SomeEnum.second_value,
349            ])
350        json_message = encoding.MessageToJson(msg)
351        self.assertEqual('{"repeated_enum": ["wire_name", "second_value"]}',
352                         json_message)
353        self.assertEqual(
354            msg, encoding.JsonToMessage(MessageWithRemappings, json_message))
355
356    def testFieldRemapping(self):
357        msg = MessageWithRemappings(another_field='abc')
358        json_message = encoding.MessageToJson(msg)
359        self.assertEqual('{"anotherField": "abc"}', json_message)
360        self.assertEqual(
361            msg, encoding.JsonToMessage(MessageWithRemappings, json_message))
362
363    def testRepeatedFieldRemapping(self):
364        msg = MessageWithRemappings(repeated_field=['abc', 'def'])
365        json_message = encoding.MessageToJson(msg)
366        self.assertEqual('{"repeatedField": ["abc", "def"]}', json_message)
367        self.assertEqual(
368            msg, encoding.JsonToMessage(MessageWithRemappings, json_message))
369
370    def testMultipleRemapping(self):
371        msg = MessageWithRemappings(
372            double_encoding=MessageWithRemappings.SomeEnum.enum_value)
373        json_message = encoding.MessageToJson(msg)
374        self.assertEqual('{"doubleEncoding": "wire_name"}', json_message)
375        self.assertEqual(
376            msg, encoding.JsonToMessage(MessageWithRemappings, json_message))
377
378    def testRepeatedRemapping(self):
379        # Should allow remapping if the mapping remains the same.
380        encoding.AddCustomJsonEnumMapping(MessageWithRemappings.SomeEnum,
381                                          'enum_value', 'wire_name')
382        encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
383                                           'double_encoding', 'doubleEncoding')
384        encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
385                                           'another_field', 'anotherField')
386        encoding.AddCustomJsonFieldMapping(MessageWithRemappings,
387                                           'repeated_field', 'repeatedField')
388
389        # Should raise errors if the remapping changes the mapping.
390        self.assertRaises(
391            exceptions.InvalidDataError,
392            encoding.AddCustomJsonFieldMapping,
393            MessageWithRemappings, 'double_encoding', 'something_else')
394        self.assertRaises(
395            exceptions.InvalidDataError,
396            encoding.AddCustomJsonFieldMapping,
397            MessageWithRemappings, 'enum_field', 'anotherField')
398        self.assertRaises(
399            exceptions.InvalidDataError,
400            encoding.AddCustomJsonEnumMapping,
401            MessageWithRemappings.SomeEnum, 'enum_value', 'another_name')
402        self.assertRaises(
403            exceptions.InvalidDataError,
404            encoding.AddCustomJsonEnumMapping,
405            MessageWithRemappings.SomeEnum, 'second_value', 'wire_name')
406
407    def testMessageToRepr(self):
408        # Using the same string returned by MessageToRepr, with the
409        # module names fixed.
410        # pylint: disable=bad-whitespace
411        msg = SimpleMessage(field='field', repfield=['field', 'field', ],)
412        # pylint: enable=bad-whitespace
413        self.assertEqual(
414            encoding.MessageToRepr(msg),
415            r"%s.SimpleMessage(field='field',repfield=['field','field',],)" % (
416                __name__,))
417        self.assertEqual(
418            encoding.MessageToRepr(msg, no_modules=True),
419            r"SimpleMessage(field='field',repfield=['field','field',],)")
420
421    def testMessageToReprWithTime(self):
422        msg = TimeMessage(timefield=datetime.datetime(
423            2014, 7, 2, 23, 33, 25, 541000,
424            tzinfo=util.TimeZoneOffset(datetime.timedelta(0))))
425        self.assertEqual(
426            encoding.MessageToRepr(msg, multiline=True),
427            ('%s.TimeMessage(\n    '
428             'timefield=datetime.datetime(2014, 7, 2, 23, 33, 25, 541000, '
429             'tzinfo=apitools.base.protorpclite.util.TimeZoneOffset('
430             'datetime.timedelta(0))),\n)') % __name__)
431        self.assertEqual(
432            encoding.MessageToRepr(msg, multiline=True, no_modules=True),
433            'TimeMessage(\n    '
434            'timefield=datetime.datetime(2014, 7, 2, 23, 33, 25, 541000, '
435            'tzinfo=TimeZoneOffset(datetime.timedelta(0))),\n)')
436
437    def testPackageMappingsNoPackage(self):
438        this_module_name = util.get_package_for_module(__name__)
439        full_type_name = 'MessageWithEnum.ThisEnum'
440        full_key = '%s.%s' % (this_module_name, full_type_name)
441        self.assertEqual(full_key,
442                         encoding._GetTypeKey(MessageWithEnum.ThisEnum, ''))
443
444    def testPackageMappingsWithPackage(self):
445        this_module_name = util.get_package_for_module(__name__)
446        full_type_name = 'MessageWithEnum.ThisEnum'
447        full_key = '%s.%s' % (this_module_name, full_type_name)
448        this_module = sys.modules[__name__]
449        new_package = 'new_package'
450        try:
451            setattr(this_module, 'package', new_package)
452            new_key = '%s.%s' % (new_package, full_type_name)
453            self.assertEqual(
454                new_key,
455                encoding._GetTypeKey(MessageWithEnum.ThisEnum, ''))
456            self.assertEqual(
457                full_key,
458                encoding._GetTypeKey(MessageWithEnum.ThisEnum, new_package))
459        finally:
460            delattr(this_module, 'package')
461
462    def testRepeatedJsonValuesAsRepeatedProperty(self):
463        encoded_msg = '{"a": [{"one": 1}]}'
464        msg = encoding.JsonToMessage(RepeatedJsonValueMessage, encoded_msg)
465        self.assertEqual(encoded_msg, encoding.MessageToJson(msg))
466
467    def testDictToProtoMap(self):
468        dict_ = {'key': 'value'}
469
470        encoded_msg = encoding.DictToProtoMap(dict_,
471                                              AdditionalPropertiesMessage)
472        expected_msg = AdditionalPropertiesMessage()
473        expected_msg.additional_properties = [
474            AdditionalPropertiesMessage.AdditionalProperty(
475                key='key', value='value')
476        ]
477        self.assertEqual(encoded_msg, expected_msg)
478
479    def testDictToProtoMapSorted(self):
480        tuples = [('key{0:02}'.format(i), 'value') for i in range(100)]
481        dict_ = dict(tuples)
482
483        encoded_msg = encoding.DictToProtoMap(dict_,
484                                              AdditionalPropertiesMessage,
485                                              sort_items=True)
486        expected_msg = AdditionalPropertiesMessage()
487        expected_msg.additional_properties = [
488            AdditionalPropertiesMessage.AdditionalProperty(
489                key=key, value=value)
490            for key, value in tuples
491        ]
492        self.assertEqual(encoded_msg, expected_msg)
493
494    def testDictToProtoMapNumeric(self):
495        dict_ = {'key': 1}
496
497        encoded_msg = encoding.DictToProtoMap(dict_,
498                                              AdditionalIntPropertiesMessage)
499        expected_msg = AdditionalIntPropertiesMessage()
500        expected_msg.additional_properties = [
501            AdditionalIntPropertiesMessage.AdditionalProperty(
502                key='key', value=1)
503        ]
504        self.assertEqual(encoded_msg, expected_msg)
505