mirror of
https://github.com/eitchtee/WYGIWYH.git
synced 2026-05-02 21:44:19 +02:00
90 lines
2.4 KiB
Python
90 lines
2.4 KiB
Python
from unittest.mock import patch
|
|
|
|
import procrastinate
|
|
from django.db import connection
|
|
from django.test import SimpleTestCase, TransactionTestCase
|
|
from procrastinate.testing import InMemoryConnector
|
|
|
|
from apps.common.procrastinate import on_app_ready
|
|
|
|
|
|
def make_app_with_task(func):
|
|
app = procrastinate.App(connector=InMemoryConnector())
|
|
task = app.task(name="sample_task")(func)
|
|
|
|
return app, task
|
|
|
|
|
|
class ProcrastinateConnectionCleanupTests(SimpleTestCase):
|
|
def test_app_ready_closes_old_connections_around_sync_tasks(self):
|
|
calls = []
|
|
|
|
def sample_task(value):
|
|
calls.append(("task", value))
|
|
return value * 2
|
|
|
|
app, task = make_app_with_task(sample_task)
|
|
|
|
with patch(
|
|
"apps.common.procrastinate.close_old_connections",
|
|
create=True,
|
|
side_effect=lambda: calls.append(("cleanup", None)),
|
|
):
|
|
on_app_ready(app)
|
|
|
|
result = task.func(3)
|
|
|
|
self.assertEqual(result, 6)
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
("cleanup", None),
|
|
("task", 3),
|
|
("cleanup", None),
|
|
],
|
|
)
|
|
|
|
def test_app_ready_closes_old_connections_when_sync_task_raises(self):
|
|
calls = []
|
|
|
|
def sample_task():
|
|
calls.append(("task", None))
|
|
raise RuntimeError("boom")
|
|
|
|
app, task = make_app_with_task(sample_task)
|
|
|
|
with patch(
|
|
"apps.common.procrastinate.close_old_connections",
|
|
create=True,
|
|
side_effect=lambda: calls.append(("cleanup", None)),
|
|
):
|
|
on_app_ready(app)
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
task.func()
|
|
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
("cleanup", None),
|
|
("task", None),
|
|
("cleanup", None),
|
|
],
|
|
)
|
|
|
|
|
|
class ProcrastinateConnectionRecoveryTests(TransactionTestCase):
|
|
def test_wrapped_task_recovers_from_closed_django_connection(self):
|
|
def sample_task():
|
|
with connection.cursor() as cursor:
|
|
cursor.execute("SELECT 1")
|
|
return cursor.fetchone()[0]
|
|
|
|
app, task = make_app_with_task(sample_task)
|
|
on_app_ready(app)
|
|
|
|
connection.ensure_connection()
|
|
connection.connection.close()
|
|
|
|
self.assertEqual(task.func(), 1)
|