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