385 lines
14 KiB
Python
385 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""Scheduler 单元测试
|
||
|
||
覆盖：
|
||
- calculate_next_run：各种调度类型的下次执行时间计算
|
||
- _parse_simple_cron：简单 cron 表达式解析
|
||
- check_and_enqueue：到期检查与入队逻辑
|
||
- start / stop：后台循环生命周期
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
from datetime import datetime, timedelta, timezone
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
from app.schemas.schedules import ScheduleConfigSchema
|
||
from app.schemas.tasks import TaskConfigSchema
|
||
from app.services.scheduler import (
|
||
Scheduler,
|
||
calculate_next_run,
|
||
_parse_simple_cron,
|
||
_parse_time,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.fixture
|
||
def sched() -> Scheduler:
|
||
return Scheduler()
|
||
|
||
|
||
@pytest.fixture
|
||
def now() -> datetime:
|
||
"""固定时间点：2025-06-10 10:00:00 UTC（周二）"""
|
||
return datetime(2025, 6, 10, 10, 0, 0, tzinfo=timezone.utc)
|
||
|
||
|
||
def _mock_cursor(fetchone_val=None, fetchall_val=None, rowcount=1):
|
||
cur = MagicMock()
|
||
cur.fetchone.return_value = fetchone_val
|
||
cur.fetchall.return_value = fetchall_val or []
|
||
cur.rowcount = rowcount
|
||
cur.__enter__ = MagicMock(return_value=cur)
|
||
cur.__exit__ = MagicMock(return_value=False)
|
||
return cur
|
||
|
||
|
||
def _mock_conn(cursor):
|
||
conn = MagicMock()
|
||
conn.cursor.return_value = cursor
|
||
return conn
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _parse_time
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestParseTime:
|
||
def test_standard_format(self):
|
||
assert _parse_time("04:00") == (4, 0)
|
||
|
||
def test_with_minutes(self):
|
||
assert _parse_time("23:45") == (23, 45)
|
||
|
||
def test_midnight(self):
|
||
assert _parse_time("00:00") == (0, 0)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# calculate_next_run — once
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNextRunOnce:
|
||
def test_once_returns_none(self, now):
|
||
cfg = ScheduleConfigSchema(schedule_type="once")
|
||
assert calculate_next_run(cfg, now) is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# calculate_next_run — interval
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNextRunInterval:
|
||
def test_interval_minutes(self, now):
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="interval", interval_value=15, interval_unit="minutes",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
assert result == now + timedelta(minutes=15)
|
||
|
||
def test_interval_hours(self, now):
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="interval", interval_value=2, interval_unit="hours",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
assert result == now + timedelta(hours=2)
|
||
|
||
def test_interval_days(self, now):
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="interval", interval_value=3, interval_unit="days",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
assert result == now + timedelta(days=3)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# calculate_next_run — daily
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNextRunDaily:
|
||
def test_daily_next_day(self, now):
|
||
cfg = ScheduleConfigSchema(schedule_type="daily", daily_time="04:00")
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 11, 4, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_daily_custom_time(self, now):
|
||
cfg = ScheduleConfigSchema(schedule_type="daily", daily_time="18:30")
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 11, 18, 30, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# calculate_next_run — weekly
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNextRunWeekly:
|
||
def test_weekly_later_this_week(self, now):
|
||
# now 是周二(2)，weekly_days=[5] 周五 → 3 天后
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="weekly", weekly_days=[5], weekly_time="08:00",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 13, 8, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_weekly_next_week(self, now):
|
||
# now 是周二(2)，weekly_days=[1] 周一 → 下周一（6天后）
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="weekly", weekly_days=[1], weekly_time="04:00",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 16, 4, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_weekly_multiple_days_picks_next(self, now):
|
||
# now 是周二(2)，weekly_days=[1, 4, 6] → 周四(4)，2 天后
|
||
cfg = ScheduleConfigSchema(
|
||
schedule_type="weekly", weekly_days=[1, 4, 6], weekly_time="09:00",
|
||
)
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 12, 9, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# calculate_next_run — cron
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNextRunCron:
|
||
def test_cron_daily(self, now):
|
||
cfg = ScheduleConfigSchema(schedule_type="cron", cron_expression="30 4 * * *")
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 11, 4, 30, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_cron_with_dow(self, now):
|
||
# "0 8 * * 5" → 每周五 08:00，now 是周二 → 周五（3天后）
|
||
cfg = ScheduleConfigSchema(schedule_type="cron", cron_expression="0 8 * * 5")
|
||
result = calculate_next_run(cfg, now)
|
||
expected = datetime(2025, 6, 13, 8, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _parse_simple_cron
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestParseSimpleCron:
|
||
def test_daily_cron(self, now):
|
||
result = _parse_simple_cron("0 4 * * *", now)
|
||
expected = datetime(2025, 6, 11, 4, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_invalid_field_count_fallback(self, now):
|
||
# 字段数不对，回退到明天 04:00
|
||
result = _parse_simple_cron("0 4 *", now)
|
||
expected = datetime(2025, 6, 11, 4, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_wildcard_hour_minute(self, now):
|
||
# "* * * * *" → hour=0, minute=0，明天 00:00
|
||
result = _parse_simple_cron("* * * * *", now)
|
||
expected = datetime(2025, 6, 11, 0, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_dow_sunday(self, now):
|
||
# "0 6 * * 0" → 每周日 06:00，now 是周二 → 周日（5天后）
|
||
result = _parse_simple_cron("0 6 * * 0", now)
|
||
expected = datetime(2025, 6, 15, 6, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_dow_same_day_future_time(self):
|
||
# 周二 08:00，cron 指定周二 12:00 → 当天
|
||
now = datetime(2025, 6, 10, 8, 0, 0, tzinfo=timezone.utc)
|
||
result = _parse_simple_cron("0 12 * * 2", now)
|
||
expected = datetime(2025, 6, 10, 12, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
def test_dow_same_day_past_time(self):
|
||
# 周二 14:00，cron 指定周二 12:00 → 下周二
|
||
now = datetime(2025, 6, 10, 14, 0, 0, tzinfo=timezone.utc)
|
||
result = _parse_simple_cron("0 12 * * 2", now)
|
||
expected = datetime(2025, 6, 17, 12, 0, 0, tzinfo=timezone.utc)
|
||
assert result == expected
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# check_and_enqueue
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestCheckAndEnqueue:
|
||
@patch("app.services.scheduler.get_connection")
|
||
@patch("app.services.scheduler.task_queue")
|
||
def test_enqueues_due_tasks(self, mock_tq, mock_get_conn, sched):
|
||
"""到期任务应被入队，且更新 last_run_at / run_count / next_run_at"""
|
||
task_config = {"tasks": ["ODS_MEMBER"], "pipeline": "api_ods_dwd"}
|
||
schedule_config = {
|
||
"schedule_type": "interval",
|
||
"interval_value": 1,
|
||
"interval_unit": "hours",
|
||
}
|
||
|
||
# 第一次 cursor：SELECT 到期任务
|
||
select_cur = _mock_cursor(
|
||
fetchall_val=[
|
||
("task-uuid-1", 42, json.dumps(task_config), json.dumps(schedule_config)),
|
||
]
|
||
)
|
||
# 第二次 cursor：UPDATE
|
||
update_cur = _mock_cursor()
|
||
|
||
conn = MagicMock()
|
||
# cursor() 依次返回 select_cur 和 update_cur
|
||
conn.cursor.side_effect = [select_cur, update_cur]
|
||
mock_get_conn.return_value = conn
|
||
|
||
mock_tq.enqueue.return_value = "queue-id-1"
|
||
|
||
count = sched.check_and_enqueue()
|
||
|
||
assert count == 1
|
||
mock_tq.enqueue.assert_called_once()
|
||
# 验证 enqueue 的参数
|
||
call_args = mock_tq.enqueue.call_args
|
||
assert call_args[0][1] == 42 # site_id
|
||
assert isinstance(call_args[0][0], TaskConfigSchema)
|
||
|
||
@patch("app.services.scheduler.get_connection")
|
||
@patch("app.services.scheduler.task_queue")
|
||
def test_no_due_tasks(self, mock_tq, mock_get_conn, sched):
|
||
"""没有到期任务时，不入队"""
|
||
cur = _mock_cursor(fetchall_val=[])
|
||
conn = _mock_conn(cur)
|
||
mock_get_conn.return_value = conn
|
||
|
||
count = sched.check_and_enqueue()
|
||
|
||
assert count == 0
|
||
mock_tq.enqueue.assert_not_called()
|
||
|
||
@patch("app.services.scheduler.get_connection")
|
||
@patch("app.services.scheduler.task_queue")
|
||
def test_skips_invalid_config(self, mock_tq, mock_get_conn, sched):
|
||
"""配置反序列化失败的任务应被跳过"""
|
||
# task_config 缺少必填字段 tasks
|
||
bad_config = {"pipeline": "api_ods_dwd"}
|
||
schedule_config = {"schedule_type": "once"}
|
||
|
||
cur = _mock_cursor(
|
||
fetchall_val=[
|
||
("task-uuid-bad", 42, json.dumps(bad_config), json.dumps(schedule_config)),
|
||
]
|
||
)
|
||
conn = _mock_conn(cur)
|
||
mock_get_conn.return_value = conn
|
||
|
||
count = sched.check_and_enqueue()
|
||
|
||
assert count == 0
|
||
mock_tq.enqueue.assert_not_called()
|
||
|
||
@patch("app.services.scheduler.get_connection")
|
||
@patch("app.services.scheduler.task_queue")
|
||
def test_enqueue_failure_continues(self, mock_tq, mock_get_conn, sched):
|
||
"""入队失败时应跳过该任务，继续处理后续任务"""
|
||
task_config = {"tasks": ["ODS_MEMBER"], "pipeline": "api_ods_dwd"}
|
||
schedule_config = {"schedule_type": "once"}
|
||
|
||
cur = _mock_cursor(
|
||
fetchall_val=[
|
||
("task-1", 42, json.dumps(task_config), json.dumps(schedule_config)),
|
||
("task-2", 42, json.dumps(task_config), json.dumps(schedule_config)),
|
||
]
|
||
)
|
||
# 需要额外的 cursor 给 UPDATE 用
|
||
update_cur = _mock_cursor()
|
||
conn = MagicMock()
|
||
conn.cursor.side_effect = [cur, update_cur]
|
||
mock_get_conn.return_value = conn
|
||
|
||
# 第一次入队失败，第二次成功
|
||
mock_tq.enqueue.side_effect = [Exception("DB error"), "queue-id-2"]
|
||
|
||
count = sched.check_and_enqueue()
|
||
|
||
assert count == 1
|
||
assert mock_tq.enqueue.call_count == 2
|
||
|
||
@patch("app.services.scheduler.get_connection")
|
||
@patch("app.services.scheduler.task_queue")
|
||
def test_once_type_sets_next_run_none(self, mock_tq, mock_get_conn, sched):
|
||
"""once 类型任务入队后，next_run_at 应被设为 NULL"""
|
||
task_config = {"tasks": ["ODS_MEMBER"], "pipeline": "api_ods_dwd"}
|
||
schedule_config = {"schedule_type": "once"}
|
||
|
||
select_cur = _mock_cursor(
|
||
fetchall_val=[
|
||
("task-uuid-1", 42, json.dumps(task_config), json.dumps(schedule_config)),
|
||
]
|
||
)
|
||
update_cur = _mock_cursor()
|
||
conn = MagicMock()
|
||
conn.cursor.side_effect = [select_cur, update_cur]
|
||
mock_get_conn.return_value = conn
|
||
mock_tq.enqueue.return_value = "queue-id-1"
|
||
|
||
sched.check_and_enqueue()
|
||
|
||
# 验证 UPDATE 语句中 next_run_at 参数为 None
|
||
update_call = update_cur.__enter__().execute.call_args
|
||
# 参数元组的第一个元素是 next_run_at
|
||
assert update_call[0][1][0] is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# start / stop 生命周期
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestLifecycle:
|
||
@pytest.mark.asyncio
|
||
async def test_stop_sets_running_false(self, sched):
|
||
sched._running = True
|
||
await sched.stop()
|
||
assert sched._running is False
|
||
assert sched._loop_task is None
|
||
|
||
def test_start_creates_task(self, sched):
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
# 在事件循环中启动
|
||
async def _run():
|
||
sched.start()
|
||
assert sched._loop_task is not None
|
||
assert not sched._loop_task.done()
|
||
await sched.stop()
|
||
|
||
loop.run_until_complete(_run())
|
||
finally:
|
||
loop.close()
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_start_stop_idempotent(self, sched):
|
||
"""多次 stop 不应报错"""
|
||
await sched.stop()
|
||
await sched.stop()
|
||
assert sched._loop_task is None
|