# -*- coding: utf-8 -*- """TaskExecutor 单元测试 覆盖:子进程启动、stdout/stderr 读取、日志广播、取消、数据库记录。 使用 asyncio 测试,mock 子进程和数据库连接避免外部依赖。 """ import asyncio from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest from app.schemas.tasks import TaskConfigSchema from app.services.task_executor import TaskExecutor @pytest.fixture def executor() -> TaskExecutor: return TaskExecutor() @pytest.fixture def sample_config() -> TaskConfigSchema: return TaskConfigSchema( tasks=["ODS_MEMBER", "ODS_PAYMENT"], flow="api_ods_dwd", store_id=42, ) def _make_stream(lines: list[bytes]) -> AsyncMock: """构造一个模拟的 asyncio.StreamReader,按行返回数据。""" stream = AsyncMock() # readline 依次返回每行,最后返回 b"" 表示 EOF stream.readline = AsyncMock(side_effect=[*lines, b""]) return stream # --------------------------------------------------------------------------- # 订阅 / 取消订阅 # --------------------------------------------------------------------------- class TestSubscription: def test_subscribe_returns_queue(self, executor: TaskExecutor): q = executor.subscribe("exec-1") assert isinstance(q, asyncio.Queue) def test_subscribe_multiple(self, executor: TaskExecutor): q1 = executor.subscribe("exec-1") q2 = executor.subscribe("exec-1") assert q1 is not q2 assert len(executor._subscribers["exec-1"]) == 2 def test_unsubscribe_removes_queue(self, executor: TaskExecutor): q = executor.subscribe("exec-1") executor.unsubscribe("exec-1", q) # 最后一个订阅者移除后,键也被清理 assert "exec-1" not in executor._subscribers def test_unsubscribe_nonexistent_is_safe(self, executor: TaskExecutor): """对不存在的 execution_id 取消订阅不应报错""" q: asyncio.Queue = asyncio.Queue() executor.unsubscribe("nonexistent", q) # --------------------------------------------------------------------------- # 广播 # --------------------------------------------------------------------------- class TestBroadcast: def test_broadcast_to_subscribers(self, executor: TaskExecutor): q1 = executor.subscribe("exec-1") q2 = executor.subscribe("exec-1") executor._broadcast("exec-1", "hello") assert q1.get_nowait() == "hello" assert q2.get_nowait() == "hello" def test_broadcast_no_subscribers_is_safe(self, executor: TaskExecutor): """无订阅者时广播不应报错""" executor._broadcast("nonexistent", "hello") def test_broadcast_end_sends_none(self, executor: TaskExecutor): q = executor.subscribe("exec-1") executor._broadcast_end("exec-1") assert q.get_nowait() is None # --------------------------------------------------------------------------- # 日志缓冲区 # --------------------------------------------------------------------------- class TestLogBuffer: def test_get_logs_empty(self, executor: TaskExecutor): assert executor.get_logs("nonexistent") == [] def test_get_logs_returns_copy(self, executor: TaskExecutor): executor._log_buffers["exec-1"] = ["line1", "line2"] logs = executor.get_logs("exec-1") assert logs == ["line1", "line2"] # 修改副本不影响原始 logs.append("line3") assert len(executor._log_buffers["exec-1"]) == 2 # --------------------------------------------------------------------------- # 执行状态查询 # --------------------------------------------------------------------------- class TestRunningState: def test_is_running_false_when_no_process(self, executor: TaskExecutor): assert executor.is_running("nonexistent") is False def test_is_running_true_when_process_active(self, executor: TaskExecutor): proc = MagicMock() proc.returncode = None executor._processes["exec-1"] = proc assert executor.is_running("exec-1") is True def test_is_running_false_when_process_exited(self, executor: TaskExecutor): proc = MagicMock() proc.returncode = 0 executor._processes["exec-1"] = proc assert executor.is_running("exec-1") is False def test_get_running_ids(self, executor: TaskExecutor): running = MagicMock() running.returncode = None exited = MagicMock() exited.returncode = 0 executor._processes["a"] = running executor._processes["b"] = exited assert executor.get_running_ids() == ["a"] # --------------------------------------------------------------------------- # _read_stream # --------------------------------------------------------------------------- class TestReadStream: @pytest.mark.asyncio async def test_read_stdout_lines(self, executor: TaskExecutor): executor._log_buffers["exec-1"] = [] stream = _make_stream([b"line1\n", b"line2\n"]) collector: list[str] = [] await executor._read_stream("exec-1", stream, "stdout", collector) assert collector == ["line1", "line2"] assert executor._log_buffers["exec-1"] == [ "[stdout] line1", "[stdout] line2", ] @pytest.mark.asyncio async def test_read_stderr_lines(self, executor: TaskExecutor): executor._log_buffers["exec-1"] = [] stream = _make_stream([b"err1\n"]) collector: list[str] = [] await executor._read_stream("exec-1", stream, "stderr", collector) assert collector == ["err1"] assert executor._log_buffers["exec-1"] == ["[stderr] err1"] @pytest.mark.asyncio async def test_read_stream_none_is_safe(self, executor: TaskExecutor): """stream 为 None 时不应报错""" collector: list[str] = [] await executor._read_stream("exec-1", None, "stdout", collector) assert collector == [] @pytest.mark.asyncio async def test_broadcast_during_read(self, executor: TaskExecutor): executor._log_buffers["exec-1"] = [] q = executor.subscribe("exec-1") stream = _make_stream([b"hello\n"]) collector: list[str] = [] await executor._read_stream("exec-1", stream, "stdout", collector) assert q.get_nowait() == "[stdout] hello" # --------------------------------------------------------------------------- # execute(集成级,mock 子进程和数据库) # --------------------------------------------------------------------------- class TestExecute: @pytest.mark.asyncio @patch("app.services.task_executor.TaskExecutor._update_execution_log") @patch("app.services.task_executor.TaskExecutor._write_execution_log") @patch("asyncio.create_subprocess_exec") async def test_successful_execution( self, mock_create, mock_write_log, mock_update_log, executor: TaskExecutor, sample_config: TaskConfigSchema, ): # 模拟子进程 proc = AsyncMock() proc.returncode = None proc.stdout = _make_stream([b"processing...\n", b"done\n"]) proc.stderr = _make_stream([]) proc.wait = AsyncMock(return_value=0) # wait 调用后设置 returncode async def _wait(): proc.returncode = 0 return 0 proc.wait = _wait mock_create.return_value = proc await executor.execute(sample_config, "exec-1", site_id=42) # 验证写入了 running 状态 mock_write_log.assert_called_once() call_kwargs = mock_write_log.call_args[1] assert call_kwargs["status"] == "running" assert call_kwargs["execution_id"] == "exec-1" # 验证更新了 success 状态 mock_update_log.assert_called_once() update_kwargs = mock_update_log.call_args[1] assert update_kwargs["status"] == "success" assert update_kwargs["exit_code"] == 0 assert "processing..." in update_kwargs["output_log"] assert "done" in update_kwargs["output_log"] # 进程已从跟踪表移除 assert "exec-1" not in executor._processes @pytest.mark.asyncio @patch("app.services.task_executor.TaskExecutor._update_execution_log") @patch("app.services.task_executor.TaskExecutor._write_execution_log") @patch("asyncio.create_subprocess_exec") async def test_failed_execution( self, mock_create, mock_write_log, mock_update_log, executor: TaskExecutor, sample_config: TaskConfigSchema, ): proc = AsyncMock() proc.returncode = None proc.stdout = _make_stream([]) proc.stderr = _make_stream([b"error occurred\n"]) async def _wait(): proc.returncode = 1 return 1 proc.wait = _wait mock_create.return_value = proc await executor.execute(sample_config, "exec-2", site_id=42) update_kwargs = mock_update_log.call_args[1] assert update_kwargs["status"] == "failed" assert update_kwargs["exit_code"] == 1 assert "error occurred" in update_kwargs["error_log"] @pytest.mark.asyncio @patch("app.services.task_executor.TaskExecutor._update_execution_log") @patch("app.services.task_executor.TaskExecutor._write_execution_log") @patch("asyncio.create_subprocess_exec") async def test_exception_during_execution( self, mock_create, mock_write_log, mock_update_log, executor: TaskExecutor, sample_config: TaskConfigSchema, ): """子进程创建失败时应记录 failed 状态""" mock_create.side_effect = OSError("command not found") await executor.execute(sample_config, "exec-3", site_id=42) update_kwargs = mock_update_log.call_args[1] assert update_kwargs["status"] == "failed" @pytest.mark.asyncio @patch("app.services.task_executor.TaskExecutor._update_execution_log") @patch("app.services.task_executor.TaskExecutor._write_execution_log") @patch("asyncio.create_subprocess_exec") async def test_subscribers_notified_on_completion( self, mock_create, mock_write_log, mock_update_log, executor: TaskExecutor, sample_config: TaskConfigSchema, ): proc = AsyncMock() proc.returncode = None proc.stdout = _make_stream([b"line\n"]) proc.stderr = _make_stream([]) async def _wait(): proc.returncode = 0 return 0 proc.wait = _wait mock_create.return_value = proc q = executor.subscribe("exec-4") await executor.execute(sample_config, "exec-4", site_id=42) # 应收到日志行 + None 哨兵 messages = [] while not q.empty(): messages.append(q.get_nowait()) assert "[stdout] line" in messages assert None in messages @pytest.mark.asyncio @patch("app.services.task_executor.TaskExecutor._update_execution_log") @patch("app.services.task_executor.TaskExecutor._write_execution_log") @patch("asyncio.create_subprocess_exec") async def test_duration_ms_recorded( self, mock_create, mock_write_log, mock_update_log, executor: TaskExecutor, sample_config: TaskConfigSchema, ): proc = AsyncMock() proc.returncode = None proc.stdout = _make_stream([]) proc.stderr = _make_stream([]) async def _wait(): proc.returncode = 0 return 0 proc.wait = _wait mock_create.return_value = proc await executor.execute(sample_config, "exec-5", site_id=42) update_kwargs = mock_update_log.call_args[1] assert isinstance(update_kwargs["duration_ms"], int) assert update_kwargs["duration_ms"] >= 0 # --------------------------------------------------------------------------- # cancel # --------------------------------------------------------------------------- class TestCancel: @pytest.mark.asyncio async def test_cancel_running_process(self, executor: TaskExecutor): proc = MagicMock() proc.returncode = None proc.terminate = MagicMock() executor._processes["exec-1"] = proc result = await executor.cancel("exec-1") assert result is True proc.terminate.assert_called_once() @pytest.mark.asyncio async def test_cancel_nonexistent_returns_false(self, executor: TaskExecutor): result = await executor.cancel("nonexistent") assert result is False @pytest.mark.asyncio async def test_cancel_already_exited_returns_false(self, executor: TaskExecutor): proc = MagicMock() proc.returncode = 0 executor._processes["exec-1"] = proc result = await executor.cancel("exec-1") assert result is False @pytest.mark.asyncio async def test_cancel_process_lookup_error(self, executor: TaskExecutor): """进程已消失时 terminate 抛出 ProcessLookupError""" proc = MagicMock() proc.returncode = None proc.terminate = MagicMock(side_effect=ProcessLookupError) executor._processes["exec-1"] = proc result = await executor.cancel("exec-1") assert result is False # --------------------------------------------------------------------------- # cleanup # --------------------------------------------------------------------------- class TestCleanup: def test_cleanup_removes_buffers_and_subscribers(self, executor: TaskExecutor): executor._log_buffers["exec-1"] = ["line"] executor.subscribe("exec-1") executor.cleanup("exec-1") assert "exec-1" not in executor._log_buffers assert "exec-1" not in executor._subscribers def test_cleanup_nonexistent_is_safe(self, executor: TaskExecutor): executor.cleanup("nonexistent")