1#!/usr/bin/env python 2 3# 4# Licensed to the Apache Software Foundation (ASF) under one 5# or more contributor license agreements. See the NOTICE file 6# distributed with this work for additional information 7# regarding copyright ownership. The ASF licenses this file 8# to you under the Apache License, Version 2.0 (the 9# "License"); you may not use this file except in compliance 10# with the License. You may obtain a copy of the License at 11# 12# http://www.apache.org/licenses/LICENSE-2.0 13# 14# Unless required by applicable law or agreed to in writing, 15# software distributed under the License is distributed on an 16# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 17# KIND, either express or implied. See the License for the 18# specific language governing permissions and limitations 19# under the License. 20# 21 22from DebugProtoTest import Srv 23from DebugProtoTest.ttypes import CompactProtoTestStruct, Empty, Wrapper 24from DebugProtoTest.ttypes import ExceptionWithAMap, MutableException 25from thrift.Thrift import TFrozenDict 26from thrift.transport import TTransport 27from thrift.protocol import TBinaryProtocol, TCompactProtocol 28import collections 29import unittest 30 31 32class TestFrozenBase(unittest.TestCase): 33 def _roundtrip(self, src, dst): 34 otrans = TTransport.TMemoryBuffer() 35 optoro = self.protocol(otrans) 36 src.write(optoro) 37 itrans = TTransport.TMemoryBuffer(otrans.getvalue()) 38 iproto = self.protocol(itrans) 39 return dst.read(iproto) or dst 40 41 def test_dict_is_hashable_only_after_frozen(self): 42 d0 = {} 43 self.assertFalse(isinstance(d0, collections.Hashable)) 44 d1 = TFrozenDict(d0) 45 self.assertTrue(isinstance(d1, collections.Hashable)) 46 47 def test_struct_with_collection_fields(self): 48 pass 49 50 def test_set(self): 51 """Test that annotated set field can be serialized and deserialized""" 52 x = CompactProtoTestStruct(set_byte_map={ 53 frozenset([42, 100, -100]): 99, 54 frozenset([0]): 100, 55 frozenset([]): 0, 56 }) 57 x2 = self._roundtrip(x, CompactProtoTestStruct()) 58 self.assertEqual(x2.set_byte_map[frozenset([42, 100, -100])], 99) 59 self.assertEqual(x2.set_byte_map[frozenset([0])], 100) 60 self.assertEqual(x2.set_byte_map[frozenset([])], 0) 61 62 def test_map(self): 63 """Test that annotated map field can be serialized and deserialized""" 64 x = CompactProtoTestStruct(map_byte_map={ 65 TFrozenDict({42: 42, 100: -100}): 99, 66 TFrozenDict({0: 0}): 100, 67 TFrozenDict({}): 0, 68 }) 69 x2 = self._roundtrip(x, CompactProtoTestStruct()) 70 self.assertEqual(x2.map_byte_map[TFrozenDict({42: 42, 100: -100})], 99) 71 self.assertEqual(x2.map_byte_map[TFrozenDict({0: 0})], 100) 72 self.assertEqual(x2.map_byte_map[TFrozenDict({})], 0) 73 74 def test_list(self): 75 """Test that annotated list field can be serialized and deserialized""" 76 x = CompactProtoTestStruct(list_byte_map={ 77 (42, 100, -100): 99, 78 (0,): 100, 79 (): 0, 80 }) 81 x2 = self._roundtrip(x, CompactProtoTestStruct()) 82 self.assertEqual(x2.list_byte_map[(42, 100, -100)], 99) 83 self.assertEqual(x2.list_byte_map[(0,)], 100) 84 self.assertEqual(x2.list_byte_map[()], 0) 85 86 def test_empty_struct(self): 87 """Test that annotated empty struct can be serialized and deserialized""" 88 x = CompactProtoTestStruct(empty_struct_field=Empty()) 89 x2 = self._roundtrip(x, CompactProtoTestStruct()) 90 self.assertEqual(x2.empty_struct_field, Empty()) 91 92 def test_struct(self): 93 """Test that annotated struct can be serialized and deserialized""" 94 x = Wrapper(foo=Empty()) 95 self.assertEqual(x.foo, Empty()) 96 x2 = self._roundtrip(x, Wrapper) 97 self.assertEqual(x2.foo, Empty()) 98 99 def test_frozen_exception(self): 100 exc = ExceptionWithAMap(blah='foo') 101 with self.assertRaises(TypeError): 102 exc.blah = 'bar' 103 mutexc = MutableException(msg='foo') 104 mutexc.msg = 'bar' 105 self.assertEqual(mutexc.msg, 'bar') 106 107 def test_frozen_exception_serialization(self): 108 result = Srv.declaredExceptionMethod_result( 109 xwamap=ExceptionWithAMap(blah="error")) 110 deserialized = self._roundtrip( 111 result, Srv.declaredExceptionMethod_result()) 112 self.assertEqual(result, deserialized) 113 114 115class TestFrozen(TestFrozenBase): 116 def protocol(self, trans): 117 return TBinaryProtocol.TBinaryProtocolFactory().getProtocol(trans) 118 119 120class TestFrozenAcceleratedBinary(TestFrozenBase): 121 def protocol(self, trans): 122 return TBinaryProtocol.TBinaryProtocolAcceleratedFactory(fallback=False).getProtocol(trans) 123 124 125class TestFrozenAcceleratedCompact(TestFrozenBase): 126 def protocol(self, trans): 127 return TCompactProtocol.TCompactProtocolAcceleratedFactory(fallback=False).getProtocol(trans) 128 129 130def suite(): 131 suite = unittest.TestSuite() 132 loader = unittest.TestLoader() 133 suite.addTest(loader.loadTestsFromTestCase(TestFrozen)) 134 suite.addTest(loader.loadTestsFromTestCase(TestFrozenAcceleratedBinary)) 135 suite.addTest(loader.loadTestsFromTestCase(TestFrozenAcceleratedCompact)) 136 return suite 137 138 139if __name__ == "__main__": 140 unittest.main(defaultTest="suite", testRunner=unittest.TextTestRunner(verbosity=2)) 141