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()