OLD | NEW |
1 #!/usr/bin/env python | 1 #!/usr/bin/env python |
2 # Copyright 2016 The LUCI Authors. All rights reserved. | 2 # Copyright 2016 The LUCI Authors. All rights reserved. |
3 # Use of this source code is governed under the Apache License, Version 2.0 | 3 # Use of this source code is governed under the Apache License, Version 2.0 |
4 # that can be found in the LICENSE file. | 4 # that can be found in the LICENSE file. |
5 | 5 |
6 import json | 6 import json |
7 import os | 7 import os |
8 import sys | 8 import sys |
9 import unittest | 9 import unittest |
10 import StringIO | 10 import StringIO |
(...skipping 66 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
77 def write(self, v): | 77 def write(self, v): |
78 self._assert_not_closed() | 78 self._assert_not_closed() |
79 self.buffer.write(v) | 79 self.buffer.write(v) |
80 | 80 |
81 def close(self): | 81 def close(self): |
82 self._assert_not_closed() | 82 self._assert_not_closed() |
83 self.closed = True | 83 self.closed = True |
84 | 84 |
85 def interpret(self): | 85 def interpret(self): |
86 data = StringIO.StringIO(self.buffer.getvalue()) | 86 data = StringIO.StringIO(self.buffer.getvalue()) |
| 87 magic = data.read(len(stream.BUTLER_MAGIC)) |
| 88 if magic != stream.BUTLER_MAGIC: |
| 89 raise ValueError('Invalid magic value ([%s] != [%s])' % ( |
| 90 magic, stream.BUTLER_MAGIC)) |
87 length, _ = varint.read_uvarint(data) | 91 length, _ = varint.read_uvarint(data) |
88 header = data.read(length) | 92 header = data.read(length) |
89 return json.loads(header), data.read() | 93 return json.loads(header), data.read() |
90 | 94 |
91 class _TestStreamClient(stream.StreamClient): | 95 class _TestStreamClient(stream.StreamClient): |
92 def __init__(self, value): | 96 def __init__(self, value): |
93 super(StreamClientTestCase._TestStreamClient, self).__init__() | 97 super(StreamClientTestCase._TestStreamClient, self).__init__() |
94 self.value = value | 98 self.value = value |
95 self.last_conn = None | 99 self.last_conn = None |
96 | 100 |
(...skipping 98 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
195 conn = client.last_conn | 199 conn = client.last_conn |
196 self.assertTrue(conn.closed) | 200 self.assertTrue(conn.closed) |
197 | 201 |
198 header, data = conn.interpret() | 202 header, data = conn.interpret() |
199 self.assertEqual(header, {'name': 'mystream', 'type': 'text'}) | 203 self.assertEqual(header, {'name': 'mystream', 'type': 'text'}) |
200 self.assertEqual(data, 'Using a text stream.') | 204 self.assertEqual(data, 'Using a text stream.') |
201 | 205 |
202 | 206 |
203 if __name__ == '__main__': | 207 if __name__ == '__main__': |
204 unittest.main() | 208 unittest.main() |
OLD | NEW |