| #!/usr/bin/env python3 |
| # Copyright 2022 The Pigweed Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not |
| # use this file except in compliance with the License. You may obtain a copy of |
| # the License at |
| # |
| # https://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT |
| # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
| # License for the specific language governing permissions and limitations under |
| # the License. |
| """Unit test for proxy.py""" |
| |
| import abc |
| import asyncio |
| from struct import pack |
| import time |
| import unittest |
| |
| from pigweed.pw_rpc.internal import packet_pb2 |
| from pigweed.pw_transfer import transfer_pb2 |
| from pw_hdlc import encode |
| from pw_transfer.chunk import Chunk, ProtocolVersion |
| |
| import proxy |
| |
| |
| class MockRng(abc.ABC): |
| def __init__(self, results: list[float]): |
| self._results = results |
| |
| def uniform(self, from_val: float, to_val: float) -> float: |
| val_range = to_val - from_val |
| val = self._results.pop() |
| val *= val_range |
| val += from_val |
| return val |
| |
| |
| class ProxyTest(unittest.IsolatedAsyncioTestCase): |
| async def test_transposer_simple(self): |
| sent_packets: list[bytes] = [] |
| new_packets_event: asyncio.Event = asyncio.Event() |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| # Notify that a new packet was "sent". |
| new_packets_event.set() |
| |
| transposer = proxy.DataTransposer( |
| lambda data: append(sent_packets, data), |
| name="test", |
| rate=0.5, |
| timeout=100, |
| seed=1234567890, |
| ) |
| transposer._rng = MockRng([0.6, 0.4]) |
| await transposer.process(b'aaaaaaaaaa') |
| await transposer.process(b'bbbbbbbbbb') |
| |
| expected_packets = [b'bbbbbbbbbb', b'aaaaaaaaaa'] |
| while True: |
| # Wait for new packets with a generous timeout. |
| try: |
| await asyncio.wait_for(new_packets_event.wait(), timeout=60.0) |
| except TimeoutError: |
| self.fail( |
| f'Timeout waiting for data. Packets sent: {sent_packets}' |
| ) |
| |
| # Only assert the sent packets are corrected when we've sent the |
| # expected number. |
| if len(sent_packets) == len(expected_packets): |
| self.assertEqual(sent_packets, expected_packets) |
| return |
| |
| async def test_transposer_timeout(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| transposer = proxy.DataTransposer( |
| lambda data: append(sent_packets, data), |
| name="test", |
| rate=0.5, |
| timeout=0.100, |
| seed=1234567890, |
| ) |
| transposer._rng = MockRng([0.4, 0.6]) |
| await transposer.process(b'aaaaaaaaaa') |
| |
| # Even though this should be transposed, there is no following data so |
| # the transposer should timout and send this in-order. |
| await transposer.process(b'bbbbbbbbbb') |
| |
| # Give the transposer time to timeout. |
| await asyncio.sleep(0.5) |
| |
| self.assertEqual(sent_packets, [b'aaaaaaaaaa', b'bbbbbbbbbb']) |
| |
| async def test_server_failure(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| packets_before_failure = [1, 2, 3] |
| server_failure = proxy.ServerFailure( |
| lambda data: append(sent_packets, data), |
| name="test", |
| packets_before_failure_list=packets_before_failure.copy(), |
| start_immediately=True, |
| ) |
| |
| # After passing the list to ServerFailure, add a test for no |
| # packets dropped |
| packets_before_failure.append(5) |
| |
| packets = [ |
| b'1', |
| b'2', |
| b'3', |
| b'4', |
| b'5', |
| ] |
| |
| for num_packets in packets_before_failure: |
| sent_packets.clear() |
| for packet in packets: |
| await server_failure.process(packet) |
| self.assertEqual(len(sent_packets), num_packets) |
| server_failure.handle_event( |
| proxy.Event( |
| proxy.EventType.TRANSFER_START, |
| Chunk(ProtocolVersion.VERSION_TWO, Chunk.Type.START), |
| ) |
| ) |
| |
| async def test_server_failure_transfer_chunks_only(self): |
| sent_packets = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| packets_before_failure = [2] |
| server_failure = proxy.ServerFailure( |
| lambda data: append(sent_packets, data), |
| name="test", |
| packets_before_failure_list=packets_before_failure.copy(), |
| start_immediately=True, |
| only_consider_transfer_chunks=True, |
| ) |
| |
| transfer_chunk = _encode_rpc_frame( |
| Chunk(ProtocolVersion.VERSION_TWO, Chunk.Type.DATA, data=b'1') |
| ) |
| |
| packets = [ |
| b'1', |
| b'2', |
| transfer_chunk, # 1 |
| b'3', |
| transfer_chunk, # 2 |
| b'4', |
| b'5', |
| transfer_chunk, # Transfer chunks should be dropped starting here. |
| transfer_chunk, |
| b'6', |
| b'7', |
| transfer_chunk, |
| ] |
| |
| for packet in packets: |
| await server_failure.process(packet) |
| |
| expected_result = [ |
| b'1', |
| b'2', |
| transfer_chunk, |
| b'3', |
| transfer_chunk, |
| b'4', |
| b'5', |
| b'6', |
| b'7', |
| ] |
| self.assertEqual(sent_packets, expected_result) |
| |
| async def test_keep_drop_queue_loop(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| keep_drop_queue = proxy.KeepDropQueue( |
| lambda data: append(sent_packets, data), |
| name="test", |
| keep_drop_queue=[2, 1, 3], |
| ) |
| |
| expected_sequence = [ |
| b'1', |
| b'2', |
| b'4', |
| b'5', |
| b'6', |
| b'9', |
| ] |
| input_packets = [ |
| b'1', |
| b'2', |
| b'3', |
| b'4', |
| b'5', |
| b'6', |
| b'7', |
| b'8', |
| b'9', |
| ] |
| |
| for packet in input_packets: |
| await keep_drop_queue.process(packet) |
| self.assertEqual(sent_packets, expected_sequence) |
| |
| async def test_keep_drop_queue(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| keep_drop_queue = proxy.KeepDropQueue( |
| lambda data: append(sent_packets, data), |
| name="test", |
| keep_drop_queue=[2, 1, 1, -1], |
| ) |
| |
| expected_sequence = [ |
| b'1', |
| b'2', |
| b'4', |
| ] |
| input_packets = [ |
| b'1', |
| b'2', |
| b'3', |
| b'4', |
| b'5', |
| b'6', |
| b'7', |
| b'8', |
| b'9', |
| ] |
| |
| for packet in input_packets: |
| await keep_drop_queue.process(packet) |
| self.assertEqual(sent_packets, expected_sequence) |
| |
| async def test_keep_drop_queue_transfer_chunks_only(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| keep_drop_queue = proxy.KeepDropQueue( |
| lambda data: append(sent_packets, data), |
| name="test", |
| keep_drop_queue=[2, 1, 1, -1], |
| only_consider_transfer_chunks=True, |
| ) |
| |
| transfer_chunk = _encode_rpc_frame( |
| Chunk(ProtocolVersion.VERSION_TWO, Chunk.Type.DATA, data=b'1') |
| ) |
| |
| expected_sequence = [ |
| b'1', |
| transfer_chunk, |
| b'2', |
| transfer_chunk, |
| b'3', |
| b'4', |
| b'5', |
| b'6', |
| b'7', |
| transfer_chunk, |
| b'8', |
| b'9', |
| b'10', |
| ] |
| input_packets = [ |
| b'1', |
| transfer_chunk, # keep |
| b'2', |
| transfer_chunk, # keep |
| b'3', |
| b'4', |
| b'5', |
| transfer_chunk, # drop |
| b'6', |
| b'7', |
| transfer_chunk, # keep |
| transfer_chunk, # drop |
| b'8', |
| transfer_chunk, # drop |
| b'9', |
| transfer_chunk, # drop |
| transfer_chunk, # drop |
| b'10', |
| ] |
| |
| for packet in input_packets: |
| await keep_drop_queue.process(packet) |
| self.assertEqual(sent_packets, expected_sequence) |
| |
| async def test_window_packet_dropper(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| window_packet_dropper = proxy.WindowPacketDropper( |
| lambda data: append(sent_packets, data), |
| name="test", |
| window_packet_to_drop=0, |
| ) |
| |
| packets = [ |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'1', |
| session_id=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'2', |
| session_id=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'3', |
| session_id=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'4', |
| session_id=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'5', |
| session_id=1, |
| ) |
| ), |
| ] |
| |
| expected_packets = packets[1:] |
| |
| # Test each even twice to assure the filter does not have issues |
| # on new window bondaries. |
| events = [ |
| proxy.Event( |
| proxy.EventType.PARAMETERS_RETRANSMIT, |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.PARAMETERS_RETRANSMIT, |
| ), |
| ), |
| proxy.Event( |
| proxy.EventType.PARAMETERS_CONTINUE, |
| Chunk( |
| ProtocolVersion.VERSION_TWO, Chunk.Type.PARAMETERS_CONTINUE |
| ), |
| ), |
| proxy.Event( |
| proxy.EventType.PARAMETERS_RETRANSMIT, |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.PARAMETERS_RETRANSMIT, |
| ), |
| ), |
| proxy.Event( |
| proxy.EventType.PARAMETERS_CONTINUE, |
| Chunk( |
| ProtocolVersion.VERSION_TWO, Chunk.Type.PARAMETERS_CONTINUE |
| ), |
| ), |
| ] |
| |
| for event in events: |
| sent_packets.clear() |
| for packet in packets: |
| await window_packet_dropper.process(packet) |
| self.assertEqual(sent_packets, expected_packets) |
| window_packet_dropper.handle_event(event) |
| |
| async def test_window_packet_dropper_extra_in_flight_packets(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so DataTransposer can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| window_packet_dropper = proxy.WindowPacketDropper( |
| lambda data: append(sent_packets, data), |
| name="test", |
| window_packet_to_drop=1, |
| ) |
| |
| packets = [ |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'1', |
| offset=0, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'2', |
| offset=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'3', |
| offset=2, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'2', |
| offset=1, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'3', |
| offset=2, |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| data=b'4', |
| offset=3, |
| ) |
| ), |
| ] |
| |
| expected_packets = packets[1:] |
| |
| # Test each even twice to assure the filter does not have issues |
| # on new window bondaries. |
| events = [ |
| None, |
| proxy.Event( |
| proxy.EventType.PARAMETERS_RETRANSMIT, |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.PARAMETERS_RETRANSMIT, |
| offset=1, |
| ), |
| ), |
| None, |
| None, |
| None, |
| None, |
| ] |
| |
| for packet, event in zip(packets, events): |
| await window_packet_dropper.process(packet) |
| if event is not None: |
| window_packet_dropper.handle_event(event) |
| |
| expected_packets = [packets[0], packets[2], packets[3], packets[5]] |
| self.assertEqual(sent_packets, expected_packets) |
| |
| async def test_event_filter(self): |
| sent_packets: list[bytes] = [] |
| |
| # Async helper so EventFilter can await on it. |
| async def append(list: list[bytes], data: bytes): |
| list.append(data) |
| |
| queue = asyncio.Queue() |
| |
| event_filter = proxy.EventFilter( |
| lambda data: append(sent_packets, data), |
| name="test", |
| event_queue=queue, |
| ) |
| |
| request = packet_pb2.RpcPacket( |
| type=packet_pb2.PacketType.REQUEST, |
| channel_id=101, |
| service_id=1001, |
| method_id=100001, |
| ).SerializeToString() |
| |
| packets = [ |
| request, |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, Chunk.Type.START, session_id=1 |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| session_id=1, |
| data=b'3', |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| session_id=1, |
| data=b'3', |
| ) |
| ), |
| request, |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, Chunk.Type.START, session_id=2 |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| session_id=2, |
| data=b'4', |
| ) |
| ), |
| _encode_rpc_frame( |
| Chunk( |
| ProtocolVersion.VERSION_TWO, |
| Chunk.Type.DATA, |
| session_id=2, |
| data=b'5', |
| ) |
| ), |
| ] |
| |
| expected_events = [ |
| None, # request |
| proxy.EventType.TRANSFER_START, |
| None, # data chunk |
| None, # data chunk |
| None, # request |
| proxy.EventType.TRANSFER_START, |
| None, # data chunk |
| None, # data chunk |
| ] |
| |
| for packet, expected_event_type in zip(packets, expected_events): |
| await event_filter.process(packet) |
| try: |
| event_type = queue.get_nowait().type |
| except asyncio.QueueEmpty: |
| event_type = None |
| self.assertEqual(event_type, expected_event_type) |
| |
| |
| def _encode_rpc_frame(chunk: Chunk) -> bytes: |
| packet = packet_pb2.RpcPacket( |
| type=packet_pb2.PacketType.SERVER_STREAM, |
| channel_id=101, |
| service_id=1001, |
| method_id=100001, |
| payload=chunk.to_message().SerializeToString(), |
| ).SerializeToString() |
| return encode.ui_frame(73, packet) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
