From 4a217c7500f074bf619fc5b19be5e76f2fb9748a Mon Sep 17 00:00:00 2001 From: Eli Lipsitz Date: Mon, 16 May 2022 23:58:20 +0000 Subject: [PATCH] pw_transfer: Support reads in Python test client This commit updates the Python client used in the pw_transfer integration test to support the new "read" action, which reads data from the server and writes it to a file specified via the config. The harness does not yet make use of this functionality. The integration test continues to pass. Bug: b/232804652 Cq-Include-Trybots: luci.pigweed.try:pigweed-integration-transfer Change-Id: I62fb6f7a102084fe92b6305f4457dedfd7a015e9 Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/94901 Commit-Queue: Armando Montanez Reviewed-by: Ted Pudlik --- pw_transfer/integration_test/python_client.py | 54 ++++++++++++------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/pw_transfer/integration_test/python_client.py b/pw_transfer/integration_test/python_client.py index b838cd2830..710d093f20 100644 --- a/pw_transfer/integration_test/python_client.py +++ b/pw_transfer/integration_test/python_client.py @@ -59,13 +59,11 @@ def _main() -> int: _LOG.critical("Failed to connect to server at %s:%d", HOSTNAME, port) return 1 - # Initialize an RPC client over the socket. + # Initialize an RPC client over the socket and set up the pw_transfer manager. rpc_client = HdlcRpcClient( lambda: rpc_socket.recv(4096), [transfer_pb2], default_channels(lambda data: rpc_socket.sendall(data)), lambda data: _LOG.info("%s", str(data))) - - # Set up the pw_transfer manager and write the data. transfer_service = rpc_client.rpcs().pw.transfer.Transfer transfer_manager = pw_transfer.Manager( transfer_service, @@ -73,26 +71,42 @@ def _main() -> int: initial_response_timeout_s=config.initial_chunk_timeout_ms / 1000, max_retries=config.max_retries, ) - for action in config.transfer_actions: - # TODO(b/232804652): Add support for reading from the server. - if action.transfer_type != config_pb2.TransferAction.TransferType.WRITE_TO_SERVER: - _LOG.critical("Only writing to the server is supported") - return 1 - try: - with open(action.file_path, 'rb') as f: - data = f.read() - except: - _LOG.exception("Failed to read input file '%s'", action.file_path) - return 1 - - try: - transfer_manager.write(action.resource_id, data) - except: - _LOG.exception("Transfer failed") + # Perform the requested transfer actions. + for action in config.transfer_actions: + if action.transfer_type == config_pb2.TransferAction.TransferType.WRITE_TO_SERVER: + try: + with open(action.file_path, 'rb') as f: + data = f.read() + except: + _LOG.critical("Failed to read input file '%s'", + action.file_path) + return 1 + + try: + transfer_manager.write(action.resource_id, data) + except: + _LOG.exception("Transfer (write to server) failed") + return 1 + elif action.transfer_type == config_pb2.TransferAction.TransferType.READ_FROM_SERVER: + try: + data = transfer_manager.read(action.resource_id) + except: + _LOG.exception("Transfer (read from server) failed") + return 1 + + try: + with open(action.file_path, 'wb') as f: + f.write(data) + except: + _LOG.critical("Failed to write output file '%s'", + action.file_path) + return 1 + else: + _LOG.critical("Unknown transfer type: %d", action.transfer_type) return 1 - _LOG.info("Transfer completed successfully") + _LOG.info("All transfers completed successfully") return 0