MediaCrawler/tests/test_monitor_scheduler.py

# -*- coding: utf-8 -*-
# Copyright (c) 2025 [email protected]
#
# This file is part of MediaCrawler project.
# Repository: https://github.com/NanmiCoder/MediaCrawler/blob/main/tests/test_monitor_scheduler.py
# GitHub: https://github.com/NanmiCoder
# Licensed under NON-COMMERCIAL LEARNING LICENSE 1.1
#
# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：
# 1. 不得用于任何商业用途。
# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。
# 3. 不得进行大规模爬取或对平台造成运营干扰。
# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。
# 5. 不得用于任何非法或不当的用途。
#
# 详细许可条款请参阅项目根目录下的LICENSE文件。
# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。

"""Tests for the monitor scheduler's firing, deferral and recovery rules."""

import pytest
import pytest_asyncio
from sqlalchemy import select

from api.monitor import db as monitor_db
from api.monitor import scheduler as scheduler_module
from api.monitor.models import (
    MODE_CREATOR,
    MonitorRun,
    MonitorTarget,
    MonitorTask,
    RUN_INTERRUPTED,
    RUN_RUNNING,
    RUN_SUCCESS,
)
from api.monitor.scheduler import MonitorScheduler
from api.monitor.settings import set_cookie, set_setting
from tools.time_util import get_current_timestamp

MS_PER_MINUTE = 60_000


class FakeCrawlerManager:
    """Stands in for the global subprocess singleton."""

    def __init__(self, busy: bool = False) -> None:
        self.busy = busy

    def is_busy(self) -> bool:
        return self.busy


@pytest_asyncio.fixture
async def db(tmp_path):
    monitor_db.set_sqlite_path(tmp_path / "monitor.db")
    await monitor_db.init_db()
    async with monitor_db.get_session() as session:
        await set_cookie(session, "web_session=test")
    yield monitor_db
    await monitor_db.dispose_engine()


@pytest_asyncio.fixture
async def executed(monkeypatch):
    """Record execute_task calls instead of launching a real crawl."""
    calls: list[tuple[int, str]] = []

    async def _fake_execute(task_id: int, trigger: str = "manual"):
        calls.append((task_id, trigger))

    monkeypatch.setattr(scheduler_module, "execute_task", _fake_execute)
    return calls


async def _make_task(
    next_run_at, enabled: bool = True, interval: int = 60, platform: str = "xhs"
) -> int:
    async with monitor_db.get_session() as session:
        now = get_current_timestamp()
        task = MonitorTask(
            name="t",
            platform=platform,
            mode=MODE_CREATOR,
            enabled=enabled,
            interval_minutes=interval,
            max_notes_count=20,
            enable_comments=True,
            max_comments_count=50,
            run_timeout_seconds=3600,
            next_run_at=next_run_at,
            last_status="idle",
            created_at=now,
            updated_at=now,
        )
        session.add(task)
        await session.flush()
        session.add(
            MonitorTarget(
                task_id=task.id,
                kind=MODE_CREATOR,
                external_id="abc123",
                xsec_token="",
                xsec_source="",
                raw_value="abc123",
                label="abc123",
                enabled=True,
                created_at=now,
            )
        )
        return task.id


async def _get_task(task_id: int) -> MonitorTask:
    async with monitor_db.get_session() as session:
        return await session.get(MonitorTask, task_id)


class TestFiring:
    @pytest.mark.asyncio
    async def test_due_task_runs_and_advances(self, monkeypatch, db, executed):
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        past = get_current_timestamp() - MS_PER_MINUTE
        task_id = await _make_task(past)

        await MonitorScheduler().tick()

        assert executed == [(task_id, "scheduled")]
        task = await _get_task(task_id)
        # Fixed-delay: the next fire is measured from now, not from the missed slot.
        assert task.next_run_at > get_current_timestamp()

    @pytest.mark.asyncio
    async def test_future_task_does_not_run(self, monkeypatch, db, executed):
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        await _make_task(get_current_timestamp() + 10 * MS_PER_MINUTE)

        await MonitorScheduler().tick()

        assert executed == []

    @pytest.mark.asyncio
    async def test_disabled_task_does_not_run(self, monkeypatch, db, executed):
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        await _make_task(get_current_timestamp() - MS_PER_MINUTE, enabled=False)

        await MonitorScheduler().tick()

        assert executed == []

    @pytest.mark.asyncio
    async def test_the_cookie_gate_reads_the_tasks_own_platform(
        self, monkeypatch, db, executed
    ):
        """cookie 闸门要按任务自己的平台取。

        以前这里是 ``get_cookie(session)``（默认小红书）—— 只有小红书时看不出问题，
        接上抖音后，抖音任务会因为读的是小红书那份 cookie 而永远不被触发，且不报错。
        """
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        async with monitor_db.get_session() as session:
            await set_cookie(session, "sessionid=dy-secret", "dy")

        task_id = await _make_task(get_current_timestamp() - MS_PER_MINUTE, platform="dy")

        await MonitorScheduler().tick()

        assert executed == [(task_id, "scheduled")]

    @pytest.mark.asyncio
    async def test_cdp_mode_frees_a_task_from_the_cookie_gate(
        self, monkeypatch, db, executed
    ):
        """开着 CDP 时不该再要求先粘 cookie。

        CDP 模式下登录态来自被接管的那台浏览器，粘不粘 cookie 都由不得它 —— 不放行的话，
        选了「接管已有 Chrome」却没粘 cookie 的用户会发现任务永远不跑，而且什么错都不报。
        """
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        async with monitor_db.get_session() as session:
            await set_setting(session, "system.cdp_enabled", "true")

        task_id = await _make_task(get_current_timestamp() - MS_PER_MINUTE, platform="dy")

        await MonitorScheduler().tick()

        assert executed == [(task_id, "scheduled")]

    @pytest.mark.asyncio
    async def test_long_outage_coalesces_into_one_run(self, monkeypatch, db, executed):
        """A missed schedule fires once, not once per missed interval."""
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        # Due two days ago on a 1-hour interval.
        await _make_task(get_current_timestamp() - 48 * 60 * MS_PER_MINUTE)

        scheduler = MonitorScheduler()
        await scheduler.tick()
        await scheduler.tick()

        assert len(executed) == 1


class TestDeferral:
    @pytest.mark.asyncio
    async def test_busy_crawler_defers_without_advancing(self, monkeypatch, db, executed):
        """A manual crawl must not consume the monitor task's slot or lose it."""
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=True))
        due_at = get_current_timestamp() - MS_PER_MINUTE
        task_id = await _make_task(due_at)

        await MonitorScheduler().tick()

        assert executed == []
        task = await _get_task(task_id)
        # Still due, so the next free tick picks it up rather than skipping a cycle.
        assert task.next_run_at == due_at

    @pytest.mark.asyncio
    async def test_deferred_task_runs_once_crawler_frees_up(self, monkeypatch, db, executed):
        fake = FakeCrawlerManager(busy=True)
        monkeypatch.setattr(scheduler_module, "crawler_manager", fake)
        task_id = await _make_task(get_current_timestamp() - MS_PER_MINUTE)

        scheduler = MonitorScheduler()
        await scheduler.tick()
        assert executed == []

        fake.busy = False
        await scheduler.tick()
        assert executed == [(task_id, "scheduled")]


class TestCookieGuard:
    @pytest.mark.asyncio
    async def test_no_cookie_blocks_run_and_keeps_task_due(self, monkeypatch, db, executed):
        """Without a cookie every run would be an auth failure; skip instead."""
        monkeypatch.setattr(scheduler_module, "crawler_manager", FakeCrawlerManager(busy=False))
        async with monitor_db.get_session() as session:
            from api.monitor.settings import cookie_key, delete_setting

            await delete_setting(session, cookie_key("xhs"))

        due_at = get_current_timestamp() - MS_PER_MINUTE
        task_id = await _make_task(due_at)

        await MonitorScheduler().tick()

        assert executed == []
        task = await _get_task(task_id)
        # Left due so it starts working the moment a cookie is pasted.
        assert task.next_run_at == due_at


class TestRecovery:
    @pytest.mark.asyncio
    async def test_running_runs_are_marked_interrupted(self, db):
        """A run left 'running' cannot be alive -- its process died with the server."""
        async with monitor_db.get_session() as session:
            now = get_current_timestamp()
            task = MonitorTask(
                name="t", platform="xhs", mode=MODE_CREATOR, enabled=True,
                interval_minutes=60, max_notes_count=20, enable_comments=True,
                max_comments_count=50, run_timeout_seconds=3600,
                next_run_at=now, last_status="running", created_at=now, updated_at=now,
            )
            session.add(task)
            await session.flush()
            session.add(
                MonitorRun(
                    task_id=task.id, trigger="scheduled", status=RUN_RUNNING,
                    phase=MODE_CREATOR, save_data_path="", queued_at=now, not_before=0,
                    started_at=now, max_comments_count=50,
                )
            )

        await MonitorScheduler().recover()

        async with monitor_db.get_session() as session:
            run = await session.scalar(select(MonitorRun))
            assert run.status == RUN_INTERRUPTED
            assert run.finished_at is not None

    @pytest.mark.asyncio
    async def test_completed_runs_are_left_alone(self, db):
        async with monitor_db.get_session() as session:
            now = get_current_timestamp()
            task = MonitorTask(
                name="t", platform="xhs", mode=MODE_CREATOR, enabled=True,
                interval_minutes=60, max_notes_count=20, enable_comments=True,
                max_comments_count=50, run_timeout_seconds=3600,
                next_run_at=now, last_status="success", created_at=now, updated_at=now,
            )
            session.add(task)
            await session.flush()
            session.add(
                MonitorRun(
                    task_id=task.id, trigger="scheduled", status=RUN_SUCCESS,
                    phase=MODE_CREATOR, save_data_path="", queued_at=now, not_before=0,
                    started_at=now, finished_at=now, max_comments_count=50,
                )
            )

        await MonitorScheduler().recover()

        async with monitor_db.get_session() as session:
            run = await session.scalar(select(MonitorRun))
            assert run.status == RUN_SUCCESS