def test_chunked_stream(self):
original_bytes = bytearray(range(100))
for data_length in [1, 10, 100]:
for buffer_length in [1, 2, 3, 5, 20, 99, 100, 101, 500]:
dest = ByteArrayOutput()
stream_out = serializers.ChunkedStream(dest, buffer_length)
stream_out.write(original_bytes[:data_length])
stream_out.close()
num_chunks = int(math.ceil(float(data_length) / buffer_length))
# length for each chunk, and a final -1 at the very end
exp_size = (num_chunks + 1) * 4 + data_length
self.assertEqual(len(dest.buffer), exp_size)
dest_pos = 0
data_pos = 0
for chunk_idx in range(num_chunks):
chunk_length = read_int(dest.buffer[dest_pos:(dest_pos +
4)])
if chunk_idx == num_chunks - 1:
exp_length = data_length % buffer_length
if exp_length == 0:
exp_length = buffer_length
else:
exp_length = buffer_length
self.assertEqual(chunk_length, exp_length)
dest_pos += 4
dest_chunk = dest.buffer[dest_pos:dest_pos + chunk_length]
orig_chunk = original_bytes[data_pos:data_pos +
chunk_length]
self.assertEqual(dest_chunk, orig_chunk)
dest_pos += chunk_length
data_pos += chunk_length
# ends with a -1
self.assertEqual(dest.buffer[-4:], write_int(-1))
def test_chunked_stream_flush(self):
wrapped = ByteArrayOutput()
stream = serializers.ChunkedStream(wrapped, 10)
stream.write(bytearray([0]))
self.assertEqual(len(wrapped.buffer), 0,
"small write should be buffered")
stream.flush()
# Expect buffer size 4 bytes + buffer data 1 byte.
self.assertEqual(len(wrapped.buffer), 5, "flush should work")
stream.close()