from __future__ import annotations

import unittest
from unittest.mock import ANY, MagicMock, patch

from govoplan_core.celery_app import (
    celery,
    dispatch_dataflow_runs,
    purge_dataflow_runs,
)


class DataflowRunWorkerTests(unittest.TestCase):
    def test_dispatch_commits_worker_outcome(self) -> None:
        session = MagicMock()
        database = MagicMock()
        database.SessionLocal.return_value.__enter__.return_value = session
        provider = MagicMock()
        provider.dispatch_pending.return_value = {
            "claimed": 1,
            "succeeded": 1,
        }

        with (
            patch(
                "govoplan_core.celery_app._dataflow_run_worker",
                return_value=provider,
            ),
            patch(
                "govoplan_core.db.session.get_database",
                return_value=database,
            ),
        ):
            result = dispatch_dataflow_runs.run(7)

        provider.dispatch_pending.assert_called_once_with(
            session,
            limit=7,
            worker_id=ANY,
        )
        session.commit.assert_called_once_with()
        self.assertEqual(1, result["succeeded"])

    def test_retention_commits_worker_outcome(self) -> None:
        session = MagicMock()
        database = MagicMock()
        database.SessionLocal.return_value.__enter__.return_value = session
        provider = MagicMock()
        provider.purge_expired.return_value = {"purged": 2}

        with (
            patch(
                "govoplan_core.celery_app._dataflow_run_worker",
                return_value=provider,
            ),
            patch(
                "govoplan_core.db.session.get_database",
                return_value=database,
            ),
        ):
            result = purge_dataflow_runs.run(25)

        provider.purge_expired.assert_called_once_with(session, limit=25)
        session.commit.assert_called_once_with()
        self.assertEqual(2, result["purged"])

    def test_routes_and_periodic_jobs_are_registered(self) -> None:
        self.assertEqual(
            celery.conf.task_routes["govoplan.dataflow.dispatch_runs"],
            {"queue": "dataflow"},
        )
        self.assertEqual(
            celery.conf.task_routes["govoplan.dataflow.purge_runs"],
            {"queue": "dataflow"},
        )
        self.assertEqual(
            "govoplan.dataflow.dispatch_runs",
            celery.conf.beat_schedule[
                "dataflow-runs-every-five-seconds"
            ]["task"],
        )
        self.assertEqual(
            "govoplan.dataflow.purge_runs",
            celery.conf.beat_schedule[
                "dataflow-run-retention-daily"
            ]["task"],
        )


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