external_demand/tests/test_same_day_retry.py

"""当天失败自动重试判定与循环逻辑测试。"""

from __future__ import annotations

import unittest
from datetime import date
from unittest.mock import patch

from app.scheduler import (
    _festival_demand_needs_retry,
    _gap_script_demand_needs_retry,
    _is_write_success,
    _run_with_same_day_retry,
    _should_run_immediately_on_startup,
)


class SameDayRetryDecisionTest(unittest.TestCase):
    def test_write_success_when_mysql_saved(self) -> None:
        self.assertTrue(
            _is_write_success({"mysql_save": {"skipped": False, "saved_count": 3}})
        )

    def test_write_success_when_partition_exists(self) -> None:
        self.assertTrue(_is_write_success({"skip_reason": "partition_data_exists"}))

    def test_festival_retry_on_missing_source(self) -> None:
        self.assertTrue(_festival_demand_needs_retry({"skip_reason": "no_demand_names"}))

    def test_festival_no_retry_when_no_active_festivals(self) -> None:
        self.assertFalse(
            _festival_demand_needs_retry({"skip_reason": "no_active_festivals"})
        )

    def test_festival_no_retry_when_nothing_to_write(self) -> None:
        self.assertFalse(_festival_demand_needs_retry({"mysql_save": {}}))

    def test_gap_retry_on_missing_source(self) -> None:
        self.assertTrue(_gap_script_demand_needs_retry({"skip_reason": "no_demand_names"}))

    def test_gap_no_retry_when_no_matched(self) -> None:
        self.assertFalse(
            _gap_script_demand_needs_retry({"skip_reason": "no_matched_demands"})
        )


class SameDayRetryLoopTest(unittest.TestCase):
    def test_retries_until_write_success(self) -> None:
        calls = {"n": 0}
        summaries = [
            {"skip_reason": "no_demand_names", "mysql_save": {}},
            {"mysql_save": {"skipped": False, "saved_count": 1}},
        ]

        def run_once_fn() -> dict:
            idx = calls["n"]
            calls["n"] += 1
            return summaries[idx]

        printed: list[dict] = []

        with patch("app.scheduler.time.sleep") as sleep_mock, patch(
            "app.scheduler._today_shanghai",
            return_value=date(2026, 8, 12),
        ):
            _run_with_same_day_retry(
                job_name="festival_demand",
                query_date=date(2026, 8, 12),
                run_once_fn=run_once_fn,
                needs_retry=_festival_demand_needs_retry,
                print_summary=printed.append,
                interval_seconds=600,
            )

        self.assertEqual(calls["n"], 2)
        sleep_mock.assert_called_once_with(600)

    def test_retries_on_exception_then_succeeds(self) -> None:
        calls = {"n": 0}

        def run_once_fn() -> dict:
            calls["n"] += 1
            if calls["n"] == 1:
                raise RuntimeError("boom")
            return {"mysql_save": {"skipped": False, "saved_count": 2}}

        with patch("app.scheduler.time.sleep") as sleep_mock, patch(
            "app.scheduler._today_shanghai",
            return_value=date(2026, 8, 12),
        ):
            _run_with_same_day_retry(
                job_name="gap_script_demand",
                query_date=date(2026, 8, 12),
                run_once_fn=run_once_fn,
                needs_retry=_gap_script_demand_needs_retry,
                print_summary=lambda _summary: None,
                interval_seconds=600,
            )

        self.assertEqual(calls["n"], 2)
        sleep_mock.assert_called_once_with(600)

    def test_stops_when_day_changes(self) -> None:
        calls = {"n": 0}
        # 第一次判定仍是当天 → sleep；醒来后跨日 → 停止
        today_values = [date(2026, 8, 12), date(2026, 8, 13)]

        def run_once_fn() -> dict:
            calls["n"] += 1
            return {"skip_reason": "no_demand_names", "mysql_save": {}}

        with patch("app.scheduler.time.sleep") as sleep_mock, patch(
            "app.scheduler._today_shanghai",
            side_effect=today_values,
        ):
            _run_with_same_day_retry(
                job_name="festival_demand",
                query_date=date(2026, 8, 12),
                run_once_fn=run_once_fn,
                needs_retry=_festival_demand_needs_retry,
                print_summary=lambda _summary: None,
                interval_seconds=600,
            )

        self.assertEqual(calls["n"], 1)
        sleep_mock.assert_called_once_with(600)


class StartupImmediateRunTest(unittest.TestCase):
    def test_skip_startup_when_today_already_succeeded(self) -> None:
        with patch(
            "app.scheduler._has_today_output_data",
            return_value=True,
        ), patch(
            "app.scheduler._today_shanghai",
            return_value=date(2026, 8, 12),
        ):
            should_run = _should_run_immediately_on_startup(
                strategy="去年同期阳历-节点事件",
                repository_cls=object,
                job_name="festival_demand",
            )
        self.assertFalse(should_run)

    def test_startup_run_when_today_not_succeeded(self) -> None:
        with patch(
            "app.scheduler._has_today_output_data",
            return_value=False,
        ), patch(
            "app.scheduler._today_shanghai",
            return_value=date(2026, 8, 12),
        ):
            should_run = _should_run_immediately_on_startup(
                strategy="当下供需gap-脚本主驱",
                repository_cls=object,
                job_name="gap_script_demand",
            )
        self.assertTrue(should_run)

    def test_startup_run_when_check_fails(self) -> None:
        with patch(
            "app.scheduler._has_today_output_data",
            side_effect=RuntimeError("db down"),
        ):
            should_run = _should_run_immediately_on_startup(
                strategy="当下供需gap-脚本主驱",
                repository_cls=object,
                job_name="gap_script_demand",
            )
        self.assertTrue(should_run)


if __name__ == "__main__":
    unittest.main()