168 lines
6.4 KiB
Python
168 lines
6.4 KiB
Python
"""
|
|
Tests for sanguo_orchestrator.runner module
|
|
Tests Orchestrator task coordination
|
|
"""
|
|
import pytest
|
|
from unittest.mock import Mock, AsyncMock, patch
|
|
from sanguo_orchestrator.runner import Orchestrator
|
|
from sanguo_orchestrator.task import TaskState
|
|
|
|
|
|
class TestOrchestrator:
|
|
"""Test Orchestrator initialization and task management"""
|
|
|
|
def test_orchestrator_initialization(self):
|
|
"""Test Orchestrator initializes correctly"""
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
assert orchestrator.db_path == "test.db"
|
|
assert orchestrator.pool.max_workers == 2
|
|
assert orchestrator._pending == {}
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
async def test_submit_cta_creates_task(self, mock_pool_class):
|
|
"""Test async submit_cta() creates task and stores spec per task_id"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
mock_pool.submit.return_value = Mock(task_id="test_1")
|
|
mock_pool.submit_work.return_value = Mock()
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
strategy_class = Mock
|
|
symbol = "AAPL"
|
|
params = {"param1": "value1"}
|
|
start = "2024-01-01"
|
|
end = "2024-12-31"
|
|
cfg = Mock()
|
|
|
|
task_id = await orchestrator.submit_cta(strategy_class, symbol, params, start, end, cfg)
|
|
|
|
# Verify task was submitted to pool
|
|
mock_pool.submit.assert_called_once()
|
|
call_args = mock_pool.submit.call_args
|
|
assert call_args[0][0] == task_id # task_id
|
|
assert call_args[0][1] == "cta" # task_type
|
|
|
|
# Verify pending spec was stored per task_id (not global)
|
|
assert task_id in orchestrator._pending
|
|
assert orchestrator._pending[task_id]["strategy_class"] == strategy_class
|
|
assert orchestrator._pending[task_id]["symbol"] == symbol
|
|
assert orchestrator._pending[task_id]["params"] == params
|
|
assert orchestrator._pending[task_id]["start"] == start
|
|
assert orchestrator._pending[task_id]["end"] == end
|
|
assert orchestrator._pending[task_id]["cfg"] == cfg
|
|
|
|
assert task_id.startswith("cta_AAPL_")
|
|
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
def test_get_status(self, mock_pool_class):
|
|
"""Test get_status() delegates to pool"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
mock_pool.get_status.return_value = TaskState.PENDING
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
status = orchestrator.get_status("test_1")
|
|
|
|
mock_pool.get_status.assert_called_once_with("test_1")
|
|
assert status == TaskState.PENDING
|
|
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
def test_get_status_nonexistent(self, mock_pool_class):
|
|
"""Test get_status() returns None for non-existent task"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
mock_pool.get_status.return_value = None
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
status = orchestrator.get_status("nonexistent")
|
|
|
|
assert status is None
|
|
|
|
@patch('sanguo_backtest.result_store.load_result')
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
def test_get_result_done_task(self, mock_pool_class, mock_load_result):
|
|
"""Test get_result() returns result for DONE task"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
|
|
mock_task = Mock()
|
|
mock_task.status = TaskState.DONE
|
|
mock_task.result_id = 12345
|
|
mock_pool.get_task.return_value = mock_task
|
|
|
|
mock_result = Mock()
|
|
mock_load_result.return_value = mock_result
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
result = orchestrator.get_result("test_1")
|
|
|
|
mock_load_result.assert_called_once_with(12345, "test.db")
|
|
assert result == mock_result
|
|
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
def test_get_result_pending_task(self, mock_pool_class):
|
|
"""Test get_result() returns None for PENDING task"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
|
|
mock_task = Mock()
|
|
mock_task.status = TaskState.PENDING
|
|
mock_pool.get_task.return_value = mock_task
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
result = orchestrator.get_result("test_1")
|
|
|
|
assert result is None
|
|
|
|
@patch('sanguo_orchestrator.runner.TaskPool')
|
|
def test_get_result_nonexistent_task(self, mock_pool_class):
|
|
"""Test get_result() returns None for non-existent task"""
|
|
mock_pool = Mock()
|
|
mock_pool_class.return_value = mock_pool
|
|
mock_pool.get_task.return_value = None
|
|
|
|
orchestrator = Orchestrator(db_path="test.db", max_workers=2)
|
|
result = orchestrator.get_result("nonexistent")
|
|
|
|
assert result is None
|
|
|
|
|
|
class TestOrchestratorAsync:
|
|
"""Test async orchestrator submit and on_stage callback"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_cta_returns_task_id_and_submits(self):
|
|
"""Test async submit_cta() returns task_id and submits to pool"""
|
|
orchestrator = Orchestrator(db_path="/tmp/t.db")
|
|
orchestrator.pool.executor = Mock()
|
|
mock_future = Mock()
|
|
orchestrator.pool.executor.submit.return_value = mock_future
|
|
|
|
with patch("sanguo_backtest.cta_engine.run_cta_backtest"):
|
|
task_id = await orchestrator.submit_cta(
|
|
Mock(), "600000", {}, "2024-01-01", "2024-06-30", cfg=Mock()
|
|
)
|
|
|
|
assert task_id.startswith("cta_600000")
|
|
orchestrator.pool.executor.submit.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_set_on_stage_callback_called(self):
|
|
"""Test set_on_stage() callback fires on task completion"""
|
|
from sanguo_orchestrator.runner import Orchestrator
|
|
from sanguo_orchestrator.task import TaskState
|
|
|
|
orchestrator = Orchestrator(db_path="/tmp/t.db")
|
|
orchestrator.pool.executor = Mock()
|
|
mock_future = Mock()
|
|
orchestrator.pool.executor.submit.return_value = mock_future
|
|
|
|
cb = AsyncMock()
|
|
orchestrator.set_on_stage(cb)
|
|
|
|
# Simulate task completion callback (with None-guard for unknown task)
|
|
await orchestrator._on_done("t1", {"statistics": {}})
|
|
|
|
# Callback should fire even though task t1 was never submitted
|
|
cb.assert_called_once() |