Add tests for Flight call options in Python

python/examples/flight/client.py

@@ -134,7 +134,9 @@ def main():
     client = pyarrow.flight.FlightClient.connect(host, port)
     while True:
         try:
-            list(client.do_action(pyarrow.flight.Action("healthcheck", b""), options=pyarrow.flight.FlightCallOptions(timeout=1)))
+            action = pyarrow.flight.Action("healthcheck", b"")
+            options = pyarrow.flight.FlightCallOptions(timeout=1)
+            list(client.do_action(action, options=options))
             break
         except pyarrow.ArrowIOError as e:
             if "Deadline" in str(e):

python/pyarrow/tests/test_flight.py

@@ -20,6 +20,7 @@ import base64
 import contextlib
 import socket
 import threading
+import time
 
 import pytest
 
@@ -93,6 +94,14 @@ class InvalidStreamFlightServer(flight.FlightServerBase):
         return flight.GeneratorStream(self.schema, [table1, table2])
 
 
+class SlowFlightServer(flight.FlightServerBase):
+    """A Flight server that delays its responses to test timeouts."""
+
+    def do_action(self, context, action):
+        time.sleep(0.5)
+        return iter([])
+
+
 class HttpBasicServerAuthHandler(flight.ServerAuthHandler):
     """An example implementation of HTTP basic authentication."""
 
@@ -257,6 +266,26 @@ def test_flight_invalid_generator_stream():
             client.do_get(flight.Ticket(b'')).read_all()
 
 
+def test_timeout_fires():
+    """Make sure timeouts fire on slow requests."""
+    # Do this in a separate thread so that if it fails, we don't hang
+    # the entire test process
+    with flight_server(SlowFlightServer) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        action = flight.Action("", b"")
+        options = flight.FlightCallOptions(timeout=0.2)
+        with pytest.raises(pa.ArrowIOError, match="Deadline Exceeded"):
+            list(client.do_action(action, options=options))
+
+
+def test_timeout_passes():
+    """Make sure timeouts do not fire on fast requests."""
+    with flight_server(ConstantFlightServer) as server_port:
+        client = flight.FlightClient.connect('localhost', server_port)
+        options = flight.FlightCallOptions(timeout=0.2)
+        client.do_get(flight.Ticket(b''), options=options).read_all()
+
+
 basic_auth_handler = HttpBasicServerAuthHandler(creds={
     b"test": b"p4ssw0rd",
 })