RamonZhou commented on code in PR #58396:
URL: https://github.com/apache/spark/pull/58396#discussion_r4446107942
##########
python/pyspark/sql/tests/connect/client/test_client.py:
##########
@@ -254,10 +254,69 @@ def GetStatus(self, req: proto.GetStatusRequest,
metadata, timeout=None):
resp.operation_statuses.append(op_status)
return resp
+ class ExitRecordingStub:
+ """Records every RPC by name, so a test can assert which RPCs the
client sent.
+
+ Answers the RPCs the exit hook may send; any other RPC is recorded and
then fails.
+ """
+
+ def __init__(self, session_id: str, release_server_side_session_id:
str = ""):
+ self._session_id = session_id
+ # ReleaseSession carries no server-side ID when the session never
existed.
+ self._release_server_side_session_id =
release_server_side_session_id
+ self.calls = []
+
+ def Config(self, req, metadata, timeout=None):
+ self.calls.append("Config")
+ resp = proto.ConfigResponse(
+ session_id=self._session_id,
server_side_session_id="server-session-id"
+ )
+ for key in req.operation.get.keys:
+ resp.pairs.add(key=key)
+ return resp
+
+ def ExecutePlan(self, req, metadata, timeout=None):
+ self.calls.append("ExecutePlan")
+ return [
+ proto.ExecutePlanResponse(
+ session_id=self._session_id,
+ server_side_session_id="server-session-id",
+ operation_id=req.operation_id,
+ )
+ ]
+
+ def ReleaseSession(self, req, metadata, timeout=None):
+ self.calls.append("ReleaseSession")
+ return proto.ReleaseSessionResponse(
+ session_id=self._session_id,
+ server_side_session_id=self._release_server_side_session_id,
+ )
+
+ def __getattr__(self, name):
+ if name.startswith("__"):
+ raise AttributeError(name)
+
+ def unexpected_rpc(*args, **kwargs):
+ self.calls.append(name)
+ raise RuntimeError(f"Unexpected RPC: {name}")
+
+ return unexpected_rpc
+
# The _cleanup_ml_cache invocation will hang in this test (no valid spark
cluster)
# and it blocks the test process exiting because it is registered as the
atexit handler
# in `SparkConnectClient` constructor. To bypass the issue, patch the
method in the test.
- SparkConnectClient._cleanup_ml_cache = lambda _: None
+ # The real method stays reachable through __wrapped__ because the test
runner can import
+ # this module twice (as __main__ and by name), and the second import must
not capture the
+ # no-op.
+ _real_cleanup_ml_cache = getattr(
+ SparkConnectClient._cleanup_ml_cache, "__wrapped__",
SparkConnectClient._cleanup_ml_cache
+ )
+
+ def _skip_cleanup_ml_cache(self):
Review Comment:
Folded this into the shared helper
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
