class TestingPipe(Pipe):
# reads come out in the same blocks as the writes
# ...wraps LocalPipe
def __init__(self, layout: str, name: str = None, stream=None):
super().__init__(name=name, layout=layout, stream=stream)
self._pipe = LocalPipe(layout, name, stream)
self._closed = False
self.data_blocks: deque = deque()
async def write(self, data):
if self._closed:
raise PipeError("Cannot write to a closed pipe")
self.data_blocks.append(data)
def write_nowait(self, data):
if self._closed:
raise PipeError("Cannot write to a closed pipe")
self.data_blocks.append(data)
def close_interval_no_wait(self):
self.data_blocks.append(None)
async def close_interval(self):
self.data_blocks.append(None)
def consume(self, rows):
return self._pipe.consume(rows)
def reread_last(self):
self._pipe.reread_last()
@property
def end_of_interval(self):
return self._pipe.end_of_interval
async def read(self, flatten=False):
# first write a block to the pipe if there are any waiting
if len(self.data_blocks) > 0:
block = self.data_blocks.popleft()
if block is None:
await self._pipe.close_interval()
else:
await self._pipe.write(block)
elif self._closed:
await self._pipe.close()
# now return the result of the inner pipe's read
return await self._pipe.read(flatten)
def is_empty(self):
return self._pipe.is_empty()
async def close(self):
self._closed = True
async def _run():
LAYOUT = "float32_3"
LENGTH = 100
pipe = LocalPipe(layout="float32_3")
tx_data = helpers.create_data(LAYOUT, length=LENGTH)
# write a block and close off the pipe
await pipe.write(tx_data)
await pipe.close()
# read it back
rx_data = await pipe.read()
np.testing.assert_array_equal(tx_data, rx_data)
# don't consume it so you can read it agai
pipe.consume(len(rx_data[:500]))
# read it back again
rx_data = await pipe.read()
np.testing.assert_array_equal(tx_data[500:], rx_data)
# the pipe should be empty now
self.assertTrue(pipe.is_empty())
with self.assertRaises(EmptyPipeError):
await pipe.read()
