vent_agent_2.0/tests/test_main_fan_monthly_schedule.py

"""月度20日均分、每日限额、冻结清单和磁盘恢复；不连接网关或模型。"""
import asyncio
from datetime import datetime
import json
from types import SimpleNamespace
from zoneinfo import ZoneInfo

import pytest
import pytest_asyncio

import config
from extensions.main_fan import batch_service, mine_catalog, monthly_schedule

TIMEZONE = ZoneInfo('Asia/Shanghai')


def at(day, *, month=11, hour=2, microsecond=0):
    return datetime(2026, month, day, hour, tzinfo=TIMEZONE, microsecond=microsecond)


@pytest_asyncio.fixture
async def monthly(monkeypatch, tmp_path):
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    for key, value in {'MAIN_FAN_BATCH_OUTPUT_DIR': str(tmp_path / 'batches'),
        'MAIN_FAN_BATCH_SPREAD_DAYS': '20', 'MAIN_FAN_BATCH_MAX_MINES': '0',
        'MAIN_FAN_BATCH_ALLOW_ALL': 'True', 'MAIN_FAN_BATCH_DEPT_IDS': 'all',
        'MAIN_FAN_BATCH_RUN_ON_STARTUP': 'False', 'MAIN_FAN_BATCH_CRON': '0 2 1 * *',
        'MAIN_FAN_REPORT_MODE': 'crawler', 'MAIN_FAN_REPORT_API_URL': 'http://reports.invalid'}.items():
        monkeypatch.setenv(key, value)
    state = SimpleNamespace(now=at(1), refreshes=0, calls=[], controllers=[], gate=None, status='completed')
    state.snapshot = {'snapshot_id': 'a' * 32, 'mines': [
        {'dept_id': str(1000 + i), 'mine_id': f'{1000 + i:08d}', 'mine_name': f'测试矿{i}'}
        for i in range(403)]}

    class Clock(datetime):
        @classmethod
        def now(cls, tz=None):
            return state.now.astimezone(tz)

    monkeypatch.setattr(monthly_schedule, 'datetime', Clock)

    async def refresh():
        state.refreshes += 1
        await asyncio.sleep(0)
        mine_catalog._write(mine_catalog._root() / state.snapshot['snapshot_id'] / 'selected-mines.json',
                            state.snapshot)
        return state.snapshot

    class Runner:
        def __init__(self, folder, pdf, **kwargs):
            self.folder = folder

        async def run(self, request):
            state.calls.append(request.model_dump(mode='json'))
            assert mine_catalog.load(request.snapshot_id)['snapshot_id'] == 'a' * 32
            if state.gate is not None:
                await state.gate.wait()
            result = {'status': state.status}
            mine_catalog._write(self.folder / 'batch.json', result)
            return result

    monkeypatch.setattr(mine_catalog, 'refresh', refresh)
    monkeypatch.setattr(batch_service, 'BatchRunner', Runner)

    def controller():
        item = batch_service.LocalBatchController()
        state.controllers.append(item)
        return item

    async def drain():
        await asyncio.gather(*(task for item in state.controllers for task in list(item.tasks.values())))

    state.controller, state.drain = controller, drain
    yield state
    for item in state.controllers:
        await item.close()


@pytest.mark.asyncio
async def test_403_mines_are_frozen_and_processed_once_over_20_days(monthly):
    controller = monthly.controller()
    monthly.now = at(1, microsecond=850000)
    first = await controller.monthly.start(['all'])
    await monthly.drain()
    plan_path = controller.monthly.root / (first['plan_id'] + '.json')
    plan = json.loads(plan_path.read_text())
    assert [len(row['dept_ids']) for row in plan['days']] == [21] * 3 + [20] * 17
    assert first['spread_days'] == 20 and first['mine_count'] == 403
    monthly.snapshot['mines'].reverse()  # 当月后续投递不能重新按上游清单切片。
    for day in range(1, 21):
        monthly.now = at(day)
        await controller.monthly.dispatch()
        await monthly.drain()
        repeated = await controller.monthly.start(['all'])
        assert repeated['plan_id'] == first['plan_id']
        assert len(monthly.calls) == day, '同一天完成后也不能领取下一份'
    assert monthly.refreshes == 1
    assert [dept for call in monthly.calls for dept in call['dept_ids']] == [str(1000 + i) for i in range(403)]
    assert {call['snapshot_id'] for call in monthly.calls} == {'a' * 32}
    assert all(call['cutoff'] is None for call in monthly.calls), '每份在实际执行时冻结自己的监测截止时间'
    monthly.now = at(21)
    assert await controller.monthly.dispatch() is None
    assert not controller.tasks


@pytest.mark.asyncio
@pytest.mark.parametrize('count', [0, 3, 40])
async def test_empty_small_and_divisible_catalogs_have_no_invalid_or_missing_batches(monthly, count):
    monthly.snapshot['mines'] = monthly.snapshot['mines'][:count]
    controller = monthly.controller()
    receipt = await controller.monthly.start(['all'])
    for day in range(1, 21):
        monthly.now = at(day)
        await controller.monthly.dispatch()
        await monthly.drain()
    assert receipt['mine_count'] == count
    assert sum(len(call['dept_ids']) for call in monthly.calls) == count
    assert len(monthly.calls) == min(count, 20)
    assert all(call['dept_ids'] for call in monthly.calls)


@pytest.mark.asyncio
async def test_monthly_and_daily_triggers_racing_only_submit_one_portion(monthly):
    controller = monthly.controller()
    await asyncio.gather(controller.monthly.start(['all']), controller.monthly.start(['all']),
                         controller.monthly.dispatch())
    await monthly.drain()
    assert monthly.refreshes == 1 and len(monthly.calls) == 1


@pytest.mark.asyncio
async def test_registered_callbacks_share_plan_and_daily_quota(monthly, monkeypatch):
    controller = monthly.controller()
    monkeypatch.setattr('extensions.get_service', lambda _: controller)
    first = await batch_service.monthly_callback()
    await monthly.drain()
    monthly.now = at(2)
    second = await batch_service.daily_callback()
    await monthly.drain()
    assert second['batch_id'] != first['daily_batch']['batch_id']
    assert monthly.refreshes == 1 and len(monthly.calls) == 2


@pytest.mark.asyncio
async def test_restart_after_completion_does_not_fetch_catalog_or_take_next_portion(monthly):
    original = monthly.controller()
    receipt = await original.monthly.start(['all'])
    await monthly.drain()
    await original.close()
    restarted = monthly.controller()
    restarted.start_on_startup()
    await restarted.monthly_task
    assert restarted.startup_task is None
    repeated = await restarted.monthly.start(['all'])
    assert repeated['daily_batch']['batch_id'] == receipt['daily_batch']['batch_id']
    assert repeated['daily_batch']['status'] == 'completed'
    assert not restarted.tasks
    assert monthly.refreshes == 1 and len(monthly.calls) == 1


@pytest.mark.asyncio
async def test_restart_resumes_interrupted_portion_with_same_id_and_consumes_that_days_quota(monthly):
    original = monthly.controller()
    receipt = await original.monthly.start(['all'])
    identifier = receipt['daily_batch']['batch_id']
    await original.close()  # worker 尚未开始；job.json 仍必须可恢复。
    assert original.get(identifier)['status'] == 'interrupted'
    monthly.now = at(5)
    restarted = monthly.controller()
    restarted.start_on_startup()
    await restarted.monthly_task
    await monthly.drain()
    assert restarted.get(identifier)['status'] == 'completed'
    assert len(monthly.calls) == 1
    await restarted.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 1, '恢复上一份的当天不能再增加新份额'
    monthly.now = at(6)
    await restarted.monthly.dispatch()
    await monthly.drain()
    assert monthly.calls[1]['dept_ids'][0] == '1021'


@pytest.mark.asyncio
async def test_downtime_defers_missed_portions_one_per_day_without_burst(monthly):
    controller = monthly.controller()
    await controller.monthly.start(['all'])
    await monthly.drain()
    monthly.now = at(5)
    await controller.monthly.dispatch()
    await monthly.drain()
    await controller.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 2
    assert monthly.calls[1]['dept_ids'] == [str(i) for i in range(1021, 1042)]
    monthly.now = at(6)
    await controller.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 3


@pytest.mark.asyncio
async def test_slow_portion_does_not_accumulate_next_days_in_queue(monthly):
    monthly.gate = asyncio.Event()
    controller = monthly.controller()
    receipt = await controller.monthly.start(['all'])
    await asyncio.sleep(0)
    monthly.now = at(2)
    second = await controller.monthly.dispatch()
    assert second['batch_id'] == receipt['daily_batch']['batch_id']
    assert len(controller.tasks) == 1 and len(monthly.calls) == 1
    monthly.gate.set()
    await monthly.drain()
    await controller.monthly.dispatch()
    assert len(monthly.calls) == 1
    monthly.now = at(3)
    await controller.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 2


@pytest.mark.asyncio
@pytest.mark.parametrize('status', ['partial', 'failed'])
async def test_failed_mines_are_recorded_without_blocking_the_next_days_portion(monthly, status):
    monthly.status = status
    controller = monthly.controller()
    first = await controller.monthly.start(['all'])
    await monthly.drain()
    assert controller.get(first['daily_batch']['batch_id'])['status'] == status
    monthly.now = at(2)
    await controller.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 2
    assert monthly.calls[1]['dept_ids'][0] == '1021'


@pytest.mark.asyncio
async def test_prior_month_backlog_and_new_plan_share_one_daily_quota(monthly):
    controller = monthly.controller()
    await controller.monthly.start(['all'])
    await monthly.drain()
    monthly.now = at(1, month=12)
    december = await controller.monthly.start(['all'])
    await monthly.drain()
    assert december['month'] == '2026-12'
    assert monthly.calls[-1]['dept_ids'][0] == '1021', '先处理上月积压，而非叠加两个月的日份额'
    await controller.monthly.dispatch()
    await monthly.drain()
    assert len(monthly.calls) == 2


@pytest.mark.asyncio
async def test_daily_or_startup_without_monthly_trigger_does_not_create_a_plan(monthly):
    controller = monthly.controller()
    monthly.now = at(15)
    controller.start_on_startup()
    await controller.monthly_task
    assert await controller.monthly.dispatch() is None
    assert monthly.refreshes == 0 and not monthly.calls


@pytest.mark.asyncio
async def test_empty_cron_disables_startup_recovery_of_an_existing_plan(monthly, monkeypatch):
    original = monthly.controller()
    await original.monthly.start(['all'])
    await original.close()
    monkeypatch.setenv('MAIN_FAN_BATCH_CRON', '')
    restarted = monthly.controller()
    restarted.start_on_startup()
    assert restarted.monthly_task is None and restarted.startup_task is None
    assert not restarted.tasks and not monthly.calls


@pytest.mark.asyncio
async def test_explicit_scope_keeps_order_and_rejects_unknown_or_duplicate_mines(monthly):
    controller = monthly.controller()
    for depts in (['1000', '1000'], ['999999']):
        with pytest.raises(ValueError, match='部门编号重复或不在'):
            await controller.monthly.start(depts)
    receipt = await controller.monthly.start(['1005', '1001', '1003'])
    await monthly.drain()
    plan = json.loads((controller.monthly.root / (receipt['plan_id'] + '.json')).read_text())
    assert [row['dept_ids'] for row in plan['days'][:3]] == [['1005'], ['1001'], ['1003']]


@pytest.mark.asyncio
async def test_daily_limit_checked_before_plan_creation_and_again_on_dispatch(monthly, monkeypatch):
    monkeypatch.setenv('MAIN_FAN_BATCH_MAX_MINES', '20')
    controller = monthly.controller()
    with pytest.raises(ValueError, match='每日批次超过'):
        await controller.monthly.start(['all'])
    assert not list(controller.monthly.root.glob('*.json')) and not controller.tasks
    controller.max_mines = 21
    await controller.monthly.start(['all'])
    await monthly.drain()
    controller.max_mines = 20
    monthly.now = at(2)
    with pytest.raises(ValueError, match='每日批次超过'):
        await controller.monthly.dispatch()
    assert len(monthly.calls) == 1


@pytest.mark.asyncio
async def test_all_mines_permission_is_not_bypassed_by_explicit_daily_slices(monthly, monkeypatch):
    controller = monthly.controller()
    monkeypatch.setenv('MAIN_FAN_BATCH_ALLOW_ALL', 'False')
    with pytest.raises(ValueError, match='全矿研判未启用'):
        await controller.monthly.start(['all'])
    assert monthly.refreshes == 0
    monkeypatch.setenv('MAIN_FAN_BATCH_ALLOW_ALL', 'True')
    await controller.monthly.start(['all'])
    await monthly.drain()
    monthly.now = at(2)
    monkeypatch.setenv('MAIN_FAN_BATCH_ALLOW_ALL', 'False')
    with pytest.raises(ValueError, match='全矿研判未启用'):
        await controller.monthly.dispatch()
    assert len(monthly.calls) == 1


@pytest.mark.parametrize('days', [0, -1, 29])
def test_invalid_spread_days_fail_early(monthly, monkeypatch, days):
    monkeypatch.setenv('MAIN_FAN_BATCH_SPREAD_DAYS', str(days))
    with pytest.raises(ValueError, match='1～28'):
        monthly.controller()


@pytest.mark.asyncio
@pytest.mark.parametrize(('spread_days', 'cron', 'count'), [(20, '15 3 1 * *', 2),
    (1, '15 3 1 * *', 1), (20, '', 0)])
async def test_scheduler_uses_beijing_monthly_and_daily_triggers(monthly, monkeypatch, spread_days, cron, count):
    from scheduler import scheduler
    monkeypatch.setenv('MAIN_FAN_LOCAL_BATCH_ENABLED', 'True')
    monkeypatch.setenv('MAIN_FAN_BATCH_SPREAD_DAYS', str(spread_days))
    monkeypatch.setenv('MAIN_FAN_BATCH_CRON', cron)
    try:
        assert await scheduler.start() == count
        jobs = {job.id: job for job in scheduler.get_scheduler().get_jobs()}
        if count:
            job = jobs['main-fan:monthly-local']
            assert job.args == (batch_service.monthly_callback,)
            assert job.trigger.get_next_fire_time(None, at(2, hour=0)).isoformat() == '2026-12-01T03:15:00+08:00'
        if count == 2:
            job = jobs['main-fan:daily-local']
            assert job.args == (batch_service.daily_callback,)
            assert job.trigger.get_next_fire_time(None, at(2, hour=0)).isoformat() == '2026-11-02T03:15:00+08:00'
    finally:
        await scheduler.shutdown()
        await asyncio.sleep(0)


@pytest.mark.asyncio
async def test_scheduler_shutdown_cancels_plan_creation_before_closing_dependencies(monthly, monkeypatch):
    from scheduler import scheduler
    controller = monthly.controller()
    monkeypatch.setattr('extensions.get_service', lambda _: controller)
    monkeypatch.setenv('MAIN_FAN_LOCAL_BATCH_ENABLED', 'True')
    entered = asyncio.Event()

    async def waiting_refresh():
        entered.set()
        await asyncio.Event().wait()

    monkeypatch.setattr(mine_catalog, 'refresh', waiting_refresh)
    try:
        await scheduler.start()
        invocation = asyncio.create_task(scheduler._invoke_batch(batch_service.monthly_callback))
        await asyncio.wait_for(entered.wait(), timeout=3)
        await scheduler.shutdown()
        assert invocation.cancelled() and not scheduler._active
        assert not controller.tasks and not list(controller.monthly.root.glob('*.json'))
    finally:
        await scheduler.shutdown()
        await asyncio.sleep(0)