1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10#   http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18#
19
20import unittest
21import os
22
23import _import_local_thrift  # noqa
24from thrift.transport import TTransport
25
26
27class TestTFileObjectTransport(unittest.TestCase):
28
29    def test_TFileObjectTransport(self):
30        test_dir = os.path.dirname(os.path.abspath(__file__))
31        datatxt_path = os.path.join(test_dir, 'data.txt')
32        buffer = '{"soft":"thrift","version":0.13,"1":true}'
33        with open(datatxt_path, "w+") as f:
34            buf = TTransport.TFileObjectTransport(f)
35            buf.write(buffer)
36            buf.flush()
37            buf.close()
38
39        with open(datatxt_path, "rb") as f:
40            buf = TTransport.TFileObjectTransport(f)
41            value = buf.read(len(buffer)).decode('utf-8')
42            self.assertEqual(buffer, value)
43            buf.close()
44        os.remove(datatxt_path)
45
46
47class TestMemoryBuffer(unittest.TestCase):
48
49    def test_memorybuffer_write(self):
50        data = '{"1":[1,"hello"],"a":{"A":"abc"},"bool":true,"num":12345}'
51
52        buffer_w = TTransport.TMemoryBuffer()
53        buffer_w.write(data.encode('utf-8'))
54        value = buffer_w.getvalue()
55        self.assertEqual(value.decode('utf-8'), data)
56        buffer_w.close()
57
58    def test_memorybuffer_read(self):
59        data = '{"1":[1, "hello"],"a":{"A":"abc"},"bool":true,"num":12345}'
60
61        buffer_r = TTransport.TMemoryBuffer(data.encode('utf-8'))
62        value_r = buffer_r.read(len(data))
63        value = buffer_r.getvalue()
64        self.assertEqual(value.decode('utf-8'), data)
65        self.assertEqual(value_r.decode('utf-8'), data)
66        buffer_r.close()
67
68
69if __name__ == '__main__':
70    unittest.main()
71