|
27 | 27 | import uuid |
28 | 28 | import zipfile |
29 | 29 | from collections.abc import Callable, Iterable |
30 | | -from socket import socketpair |
| 30 | +from socket import socket, socketpair |
31 | 31 | from typing import TYPE_CHECKING, Any, BinaryIO |
32 | 32 | from unittest.mock import MagicMock, PropertyMock, patch |
33 | 33 |
|
|
54 | 54 | from airflow.dag_processing.dagbag import DagBag |
55 | 55 | from airflow.dag_processing.manager import DagFileProcessorManager, process_parse_results |
56 | 56 | from airflow.dag_processing.processor import ( |
| 57 | + BaseDagFileProcessorProcess, |
57 | 58 | DagFileParseRequest, |
58 | 59 | DagFileParsingResult, |
59 | 60 | DagFileProcessorProcess, |
@@ -2546,6 +2547,25 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): |
2546 | 2547 | } |
2547 | 2548 |
|
2548 | 2549 |
|
| 2550 | +class TestBaseDagFileProcessorProcess: |
| 2551 | + @patch.object(BaseDagFileProcessorProcess, "cleanup_sockets_after_kill", autospec=True) |
| 2552 | + def test_close_without_log_file_handle_cleans_up_sockets(self, cleanup_sockets_after_kill): |
| 2553 | + proc = BaseDagFileProcessorProcess( |
| 2554 | + process_log=structlog.get_logger(), |
| 2555 | + id=uuid.uuid4(), |
| 2556 | + pid=1234, |
| 2557 | + process=MagicMock(spec=supervisor.ProcessTracker), |
| 2558 | + stdin=MagicMock(spec=socket), |
| 2559 | + client=MagicMock(spec=Client), |
| 2560 | + bundle_name="mybundle", |
| 2561 | + dag_file_rel_path="dags/my_dag.py", |
| 2562 | + ) |
| 2563 | + |
| 2564 | + proc.close() |
| 2565 | + |
| 2566 | + cleanup_sockets_after_kill.assert_called_once_with(proc) |
| 2567 | + |
| 2568 | + |
2549 | 2569 | class TestMultiTeamCallbackMetrics: |
2550 | 2570 | """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" |
2551 | 2571 |
|
|
0 commit comments