henry3260 commented on code in PR #65587:
URL: https://github.com/apache/airflow/pull/65587#discussion_r4012866055
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -3782,6 +3783,130 @@ def _handle_request(self, msg, log, req_id):
assert isinstance(response, VariableResult)
assert response.value == "value"
+ @pytest.fixture
+ def socket_supervisor(self, mocker, socket_pair):
+ """An in-process supervisor wired to a socket, as
``_setup_subprocess_socket`` leaves it."""
+ read_end, write_end = socket_pair
+
+ supervisor = InProcessTestSupervisor(
+ id=TI_ID,
+ pid=12345,
+ process=mocker.Mock(),
+ process_log=mocker.MagicMock(),
+ client=mocker.MagicMock(spec=sdk_client.Client),
+ )
+ supervisor.comms = InProcessSupervisorComms(supervisor=supervisor)
+ supervisor.stdin = write_end
+ supervisor.client.variables.get.return_value =
VariableResult(key="test_key", value="test_value")
+
+ return supervisor, read_end
+
+ @patch("airflow.sdk.execution_time.request_handlers.mask_secret")
+ @pytest.mark.parametrize("req_id", [0, 42], ids=["first_request",
"later_request"])
+ def test_socket_request_is_answered_over_the_socket(
+ self, mock_mask_secret, socket_supervisor, mocker, req_id
+ ):
+ """A virtualenv operator under ``dag.test()`` runs in a real child
process that reconnects to
+ the supervisor over ``__AIRFLOW_SUPERVISOR_FD``, so its requests can
only be answered with a
+ response frame on that socket. ``req_id=0`` is deliberate: the child's
``CommsDecoder``
+ numbers its requests from 0, so the id cannot tell the two paths apart.
+ """
+ supervisor, read_end = socket_supervisor
+
+ generator = supervisor.handle_requests(log=mocker.Mock())
+ next(generator)
+ generator.send(_RequestFrame(id=req_id,
body=GetVariable(key="test_key").model_dump()))
+
+ read_end.settimeout(1)
+ frame_len = int.from_bytes(read_end.recv(4), "big")
+ frame =
msgspec.msgpack.Decoder(_ResponseFrame).decode(read_end.recv(frame_len))
+
+ assert frame.id == req_id
+ assert frame.body == {"key": "test_key", "value": "test_value",
"type": "VariableResult"}
+
+ @patch("airflow.sdk.execution_time.request_handlers.mask_secret")
+ def test_in_process_request_is_not_written_to_the_socket(self,
mock_mask_secret, socket_supervisor):
+ """The task running in this process reads its response from the queue,
not the socket."""
+ supervisor, read_end = socket_supervisor
+
+ response = supervisor.comms.send(GetVariable(key="test_key"))
+
+ assert response == VariableResult(key="test_key", value="test_value")
+ read_end.settimeout(0.1)
+ with pytest.raises(TimeoutError):
+ read_end.recv(1)
+
+ def test_concurrent_in_process_requests_get_their_own_response(self,
mocker):
+ """Requests in flight at the same time must not be answered with each
other's response.
Review Comment:
> This docstring describes the socket being serviced on its own thread while
the in-process task has a request outstanding, but there is no socket in this
test: `ConcurrentSupervisor` is built without `stdin` and both requests go
through `comms.send` as ordinary in-process callers. It is a real test of the
per-call sink, just not of the scenario named, and the interleaving that does
happen in production has no test at all: on the socket thread `mask_secret`
re-enters `comms.send` inside the child's `GetVariable`, which is the one case
where a live sink and the `sink is None` branch coexist (all three socket tests
patch `mask_secret` out). Also `first.join(10)` below has no `assert not
first.is_alive()`, so a regression that hangs the first caller surfaces 20s
later as an opaque dict comparison instead of the clear message the sibling
test gives at line 3906.
Applied!
--
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]
