SameerMesiah97 commented on code in PR #76555:
URL: https://github.com/apache/airflow/pull/76555#discussion_r4518961348
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -4695,6 +4695,47 @@ def
test_handle_requests_network_exception_does_not_crash_loop(self, watched_sub
# Should not raise StopIteration (which would mean the loop crashed).
generator.send(req2)
+ @pytest.mark.parametrize(
+ "invalid_body",
+ [
+ pytest.param({"type": "NonExistentMessageType"},
id="unknown_message_type"),
+ pytest.param({"type": "GetVariable"}, id="missing_required_field"),
+ ],
+ )
+ def test_handle_requests_decoding_error_sends_error_response(
+ self, watched_subprocess, mocker, invalid_body
+ ):
+ """Test that invalid request frames produce an error response rather
than dropping the frame."""
+ watched_subprocess, read_socket = watched_subprocess
+
+ generator = watched_subprocess.handle_requests(log=mocker.Mock())
+ next(generator)
+
+ req_frame = _RequestFrame(id=randint(1, 2**32 - 1), body=invalid_body)
+ generator.send(req_frame)
+
+ read_socket.settimeout(0.5)
+ frame_len = int.from_bytes(read_socket.recv(4), "big")
+ bytes_ = read_socket.recv(frame_len)
+ frame = msgspec.msgpack.Decoder(_ResponseFrame).decode(bytes_)
+
+ assert frame.id == req_frame.id
+ assert frame.error is not None
+ assert frame.error["error"] == "GENERIC_ERROR"
+ assert "Unable to decode message" in frame.error["detail"]["message"]
+
+ comms = CommsDecoder(socket=None)
+ with pytest.raises(AirflowRuntimeError) as exc_info:
+ comms._from_frame(frame)
+
+ assert exc_info.value.error.error == ErrorType.GENERIC_ERROR
+ assert exc_info.value.error.detail["exception_type"] is not None
+
+ watched_subprocess.client.task_instances.succeed =
mocker.Mock(return_value=None)
+ msg = SucceedTask(end_date=timezone.parse("2024-10-31T12:00:00Z"))
+ req_valid = _RequestFrame(id=randint(1, 2**32 - 1),
body=msg.model_dump())
+ generator.send(req_valid)
Review Comment:
You could add this assert:
`watched_subprocess.client.task_instances.succeed.assert_called_once()`
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1179,8 +1179,21 @@ def handle_requests(self, log: FilteringBoundLogger) ->
Generator[None, _Request
try:
msg =
self.decoder.validate_python(self._deserialize_request(request.body))
- except Exception:
+ except Exception as e:
log.exception("Unable to decode message", body=request.body)
+ with suppress(Exception):
Review Comment:
This is a bit broad but since this is intended to handle unexpected errors,
it is fine.
--
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]
