Agent/examples/auto_put_ad_mini/test_daily_roi_polling.py

import json
import sys
import unittest
from datetime import date, datetime
from pathlib import Path
from unittest.mock import Mock, patch


HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import run_daily_service  # noqa: E402
from roi_control import repository  # noqa: E402
from roi_control.config import RoiConfig  # noqa: E402
from roi_control.data_source import SHANGHAI, SourceDataNotReadyError  # noqa: E402


class RoiPollingConfigTest(unittest.TestCase):
    def test_default_polling_configuration_is_30_minutes_until_1800(self):
        config = RoiConfig()
        config.validate()

        self.assertEqual(30, config.daily_poll_interval_minutes)
        self.assertEqual((18, 0), (
            config.daily_poll_cutoff_hour,
            config.daily_poll_cutoff_minute,
        ))
        self.assertEqual("0,30", run_daily_service._roi_poll_minutes(config))

    def test_poll_interval_must_divide_one_hour(self):
        with self.assertRaisesRegex(ValueError, "positive divisor of 60"):
            RoiConfig(daily_poll_interval_minutes=25).validate()

    def test_cutoff_must_be_later_than_report_time(self):
        with self.assertRaisesRegex(ValueError, "cutoff must be later"):
            RoiConfig(
                report_hour=18,
                report_minute=0,
                daily_poll_cutoff_hour=18,
                daily_poll_cutoff_minute=0,
            ).validate()


class DailyRoiPollingTest(unittest.TestCase):
    def setUp(self):
        self.config = RoiConfig(
            daily_enabled=True,
            report_hour=9,
            report_minute=0,
            daily_poll_interval_minutes=30,
            daily_poll_cutoff_hour=18,
            daily_poll_cutoff_minute=0,
        )
        self.active_time = datetime(2026, 8, 14, 9, 30, tzinfo=SHANGHAI)
        self.cutoff_time = datetime(2026, 8, 14, 18, 0, tzinfo=SHANGHAI)

    def test_existing_formal_run_skips_source_check_and_report(self):
        existing = {"run_id": "roi-existing", "status": "PENDING_APPROVAL"}
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=existing,
            ),
            patch.object(run_daily_service, "validate_source_ready") as validate,
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.poll_daily_roi(
                self.config,
                now=self.active_time,
            )

        self.assertEqual("ALREADY_RAN", result)
        validate.assert_not_called()
        run.assert_not_called()

    def test_cutoff_skips_existing_formal_run_without_alert(self):
        existing = {"run_id": "roi-existing", "status": "COMPLETED"}
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=existing,
            ),
            patch.object(run_daily_service, "validate_source_ready") as validate,
            patch.object(run_daily_service, "_notify_roi_source_cutoff") as notify,
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.finalize_daily_roi_at_cutoff(
                self.config,
                now=self.cutoff_time,
            )

        self.assertEqual("ALREADY_RAN", result)
        validate.assert_not_called()
        notify.assert_not_called()
        run.assert_not_called()

    def test_missing_source_waits_without_failure_alert(self):
        error = SourceDataNotReadyError("required_dt=20260813")
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=None,
            ),
            patch.object(run_daily_service, "ODPSClient"),
            patch.object(
                run_daily_service,
                "validate_source_ready",
                side_effect=error,
            ),
            patch.object(run_daily_service, "_notify_roi_source_cutoff") as notify,
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.poll_daily_roi(
                self.config,
                now=self.active_time,
            )

        self.assertEqual("WAITING_SOURCE", result)
        notify.assert_not_called()
        run.assert_not_called()

    def test_ready_source_triggers_formal_run_once(self):
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=None,
            ),
            patch.object(run_daily_service, "ODPSClient"),
            patch.object(
                run_daily_service,
                "validate_source_ready",
                return_value={"row_count": 10, "self_rows": 6, "gzh_rows": 4},
            ) as validate,
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.poll_daily_roi(
                self.config,
                now=self.active_time,
            )

        self.assertEqual("TRIGGERED", result)
        validate.assert_called_once_with(unittest.mock.ANY, "20260813")
        run.assert_called_once_with()

    def test_cutoff_performs_final_check_and_alerts_when_still_missing(self):
        error = SourceDataNotReadyError("required_dt=20260813")
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=None,
            ),
            patch.object(run_daily_service, "ODPSClient"),
            patch.object(
                run_daily_service,
                "validate_source_ready",
                side_effect=error,
            ),
            patch.object(run_daily_service, "_notify_roi_source_cutoff") as notify,
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.finalize_daily_roi_at_cutoff(
                self.config,
                now=self.cutoff_time,
            )

        self.assertEqual("CUTOFF_NOT_READY", result)
        notify.assert_called_once_with(
            target_date="20260813",
            cutoff_time="18:00",
            error="required_dt=20260813",
        )
        run.assert_not_called()

    def test_cutoff_runs_when_source_becomes_ready(self):
        with (
            patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
                return_value=None,
            ),
            patch.object(run_daily_service, "ODPSClient"),
            patch.object(
                run_daily_service,
                "validate_source_ready",
                return_value={"row_count": 10, "self_rows": 6, "gzh_rows": 4},
            ),
            patch.object(run_daily_service, "run_daily_roi") as run,
        ):
            result = run_daily_service.finalize_daily_roi_at_cutoff(
                self.config,
                now=self.cutoff_time,
            )

        self.assertEqual("TRIGGERED", result)
        run.assert_called_once_with()

    def test_poll_does_nothing_outside_active_window(self):
        for now in (
            datetime(2026, 8, 14, 8, 59, tzinfo=SHANGHAI),
            self.cutoff_time,
            datetime(2026, 8, 14, 18, 30, tzinfo=SHANGHAI),
        ):
            with self.subTest(now=now), patch.object(
                run_daily_service,
                "load_formal_run_for_end_date",
            ) as load_run:
                result = run_daily_service.poll_daily_roi(self.config, now=now)
                self.assertEqual("OUTSIDE_WINDOW", result)
                load_run.assert_not_called()

    def test_scheduler_registers_poll_and_cutoff_jobs(self):
        scheduler = Mock()
        scheduler.get_jobs.return_value = []
        with (
            patch.object(
                run_daily_service.RoiConfig,
                "from_env",
                return_value=self.config,
            ),
            patch.object(run_daily_service, "validate_region_dictionary"),
            patch.object(run_daily_service, "initialize_schema"),
            patch.object(run_daily_service, "_env_flag", return_value=False),
            patch.object(
                run_daily_service,
                "BlockingScheduler",
                return_value=scheduler,
            ),
        ):
            run_daily_service.main([])

        jobs = {call.kwargs["id"]: call for call in scheduler.add_job.call_args_list}
        self.assertIn("daily_roi", jobs)
        self.assertIn("daily_roi_cutoff", jobs)
        self.assertIs(jobs["daily_roi"].args[0], run_daily_service.poll_daily_roi)
        self.assertIs(
            jobs["daily_roi_cutoff"].args[0],
            run_daily_service.finalize_daily_roi_at_cutoff,
        )


class FormalRunLookupTest(unittest.TestCase):
    def test_internal_test_is_ignored_when_formal_run_exists(self):
        cursor = Mock()
        cursor.__enter__ = Mock(return_value=cursor)
        cursor.__exit__ = Mock(return_value=False)
        cursor.fetchall.return_value = [
            {
                "run_id": "roi-internal_internal_test",
                "config_json": json.dumps({"internal_test": True}),
            },
            {
                "run_id": "roi-formal",
                "config_json": json.dumps({"internal_test": False}),
            },
        ]
        connection = Mock()
        connection.cursor.return_value = cursor

        with patch.object(repository, "connect", return_value=connection):
            result = repository.load_formal_run_for_end_date(date(2026, 8, 13))

        self.assertEqual("roi-formal", result["run_id"])
        connection.close.assert_called_once_with()


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