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 <[email protected]>
Reviewed-by: Ted Pudlik <[email protected]>

pw_transfer/integration_test/python_client.py

@@ -59,40 +59,54 @@ 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,
         default_response_timeout_s=config.chunk_timeout_ms / 1000,
         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