amoghrajesh commented on code in PR #44241:
URL: https://github.com/apache/airflow/pull/44241#discussion_r1859713779
##########
task_sdk/tests/execution_time/test_task_runner.py:
##########
@@ -60,6 +60,65 @@ def test_recv_StartupDetails(self):
assert decoder.request_socket.writable()
assert decoder.request_socket.fileno() == w2.fileno()
+ @pytest.mark.usefixtures("disable_capturing")
+ @pytest.mark.parametrize(
+ "msg",
+ [
+ DeferTask(
+ state="deferred",
+ classpath="my-class-path",
+ kwargs={},
+ next_method="execute_callback",
+ timeout=None,
+ )
+ ],
+ )
+ def test_send_message(self, msg):
+ r, w = socketpair()
+ # Create a valid FD for the decoder to open
+ _, w2 = socketpair()
+
+ w.makefile("wb").write(
+ b'{"type":"StartupDetails", "ti": {'
+ b'"id": "4d828a62-a417-4936-a7a6-2b3fabacecab", "task_id": "a",
"try_number": 1, "run_id": "b", "dag_id": "c" }, '
+ b'"file": "/dev/null", "requests_fd": ' +
str(w2.fileno()).encode("ascii") + b"}\n"
+ )
+
+ decoder = CommsDecoder(input=r.makefile("r"))
+
+ msg = decoder.get_message()
+ assert isinstance(msg, StartupDetails)
+ assert msg.ti.id == uuid.UUID("4d828a62-a417-4936-a7a6-2b3fabacecab")
+ assert msg.ti.task_id == "a"
+ assert msg.ti.dag_id == "c"
+ assert msg.file == "/dev/null"
+
+ # Since this was a StartupDetails message, the decoder should open the
other socket
+ assert decoder.request_socket is not None
+ assert decoder.request_socket.writable()
+ assert decoder.request_socket.fileno() == w2.fileno()
+
+ msg_to_send = DeferTask(
+ state="deferred",
+ kwargs={},
+ next_method="execute_callback",
+ timeout=None,
+ )
+
+ decoder.send_request(msg_to_send)
+
+ w2.sendall(msg_to_send.model_dump_json().encode() + b"\n")
+
+ data = w2.recv(1024)
+
+ expected_message = msg_to_send.model_dump_json().encode() + b"\n"
+ assert data == expected_message
+
+ decoder.send_request(log=mock.MagicMock(), msg=msg_to_send)
+ # Check if `w2` received the message
+ sent_message = w2.recv(1024) # Adjust buffer size if necessary
+ assert sent_message == b'{"model": "dumped_json"}\n' # Verify the
message content
Review Comment:
This test is a work in progress, I need to update this one to test if the
exception was raised and right message was sent
--
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]
