/
tests.py
204 lines (156 loc) · 5.98 KB
/
tests.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
import os
from dateutil import tz
from StringIO import StringIO
from datetime import datetime, timedelta
from flask import Flask, request, abort
from flask.ext.transit import init_transit, transition, register_handlers
from flask.ext.testing import TestCase
from transit.writer import Writer
from transit.reader import Reader
class TestObj(object):
def __init__(self, number):
self.number = number
def __eq__(self, other):
return self.number == other.number
class TestObj2(TestObj):
pass
class TestObjHandler(object):
@staticmethod
def tag(_):
return "testobj"
@staticmethod
def rep(test_obj):
return test_obj.number
@staticmethod
def string_rep(test_obj):
return str(test_obj.number)
@staticmethod
def from_rep(value):
return TestObj(int(value))
class TestObj2Handler(object):
@staticmethod
def tag(_):
return "testobj2"
@staticmethod
def from_rep(value):
return TestObj2(int(value))
def make_app(name):
''' Constructs a test app '''
app = Flask('tests')
# TODO: Thinking this config format for custom readers and writers may suck.
# Either switch to a dict, or combined reader & writers? Easy to do with mixins
# if we really want to seperate read and write...
init_transit(app, {TestObj: TestObjHandler})
@app.route('/echo_transit/<protocol>', methods=['POST'])
def echo_transit(protocol):
return to_transit(request.transit, protocol)
@app.route('/echo_transition/<protocol>', methods=['POST'])
def echo_transition(protocol):
return transition(request.transit, protocol)
@app.route('/expect_no_transit', methods=['POST'])
def expect_no_transit():
if request.transit:
abort(400)
return 'ok'
return app
def to_transit(in_data, protocol='json'):
io = StringIO()
writer = Writer(io, protocol)
writer.register(TestObj, TestObjHandler)
writer.write(in_data)
return io.getvalue()
def from_transit(in_data, protocol):
io = StringIO(in_data)
reader = Reader(protocol)
reader.register("testobj", TestObjHandler)
return Reader(protocol).read(io)
class FlaskTransitTests(TestCase):
def create_app(self):
app = make_app('test')
app.config['TESTING'] = True
return app
def _reading_test(self, in_data, base_url, protocol):
data = to_transit(in_data, protocol)
response = self.client.post(
os.path.join(base_url, protocol),
data=data,
headers={'content-type': 'application/transit+' + protocol}
)
self.assertEqual(response.data, data)
self.assertEqual(from_transit(response.data, protocol), in_data)
return response
def test_transit_json_reading(self):
self._reading_test({'hi': 'there',
'ls': (1, 2, 3),
'aset': frozenset({1, 2, 3})},
'/echo_transit/',
'json')
def test_transit_msgpack_reading(self):
self._reading_test({'hi': 'there',
'ls': (1, 2, 3),
'aset': frozenset({1, 2, 3})},
'/echo_transit/',
'msgpack')
def test_transition_json(self):
self._reading_test({'hi': 'there',
'ls': (1, 2, 3),
'aset': frozenset({1, 2, 3})},
'/echo_transition/',
'json')
def test_transition_msgpack(self):
self._reading_test({'hi': 'there',
'ls': (1, 2, 3),
'aset': frozenset({1, 2, 3})},
'/echo_transition/',
'msgpack')
def test_custom_rw_json(self):
self._reading_test({'hi': 'there',
'obj': TestObj(3)},
'/echo_transition/',
'json')
def test_custom_rw_msgpack(self):
self._reading_test({'hi': 'there',
'obj': TestObj(100)},
'/echo_transition/',
'msgpack')
def _do_datetime_test(self, protocol):
# datetime tests need to be done slightly differently - using
# assertAlmostEqual instead. This may be needed because datetimes are
# represented as floats under the hood? Need to check that though.
in_data = {'x': datetime.now(tz=tz.tzutc())}
data = to_transit(in_data, protocol)
response = self.client.post(
'/echo_transit/' + protocol,
data=data,
headers={'content-type': 'application/transit+' + protocol}
)
self.assertAlmostEqual(
from_transit(response.data, protocol)['x'],
in_data['x'],
delta=timedelta(milliseconds=10)
)
def test_datetime_json(self):
self._do_datetime_test('json')
def test_datetime_msgpack(self):
self._do_datetime_test('msgpack')
def test_transit_ignores_json(self):
response = self.client.post(
'/expect_no_transit',
data='{"hi": "there"}',
headers={'content-type': 'application/json'}
)
self.assertEqual(response.status_code, 200)
def test_register_handlers_adds_handlers(self):
register_handlers(self.app,
{TestObj2: TestObj2Handler})
self._reading_test({'hi': 'there',
'obj': TestObj2(100)},
'/echo_transition/',
'msgpack')
def test_register_handlers_keeps_existing_handlers(self):
register_handlers(self.app,
{TestObj2: TestObj2Handler})
self._reading_test({'hi': 'there',
'obj': TestObj(100)},
'/echo_transition/',
'msgpack')