92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
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()
|