diff --git a/tests/test_executor.py b/tests/test_executor.py index a2729eb..b7af974 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -7,6 +7,7 @@ import time import uuid +import pytest from django.tasks import ( TaskContext, TaskResult, @@ -182,6 +183,7 @@ def test_run__executes_model_task_in_spawned_worker(self): finally: multiprocessing.set_start_method(original_start_method, force=True) + @pytest.mark.django_db(transaction=True) def test_worker_acquires_updates_and_acknowledges(self): """Worker acquires, executes, and acknowledges via its own backend.""" enqueued = default_task_backend.enqueue(echo, args=[42]) @@ -288,6 +290,8 @@ def test_shutdown_requested__is_settable(self): class TestWorkerThread: """Tests for the WorkerThread class.""" + pytestmark = pytest.mark.django_db(transaction=True) + def test_execute_task_result__successful_execution(self): """execute_task_result runs a task and returns SUCCESSFUL result.""" result = WorkerThread( diff --git a/tests/test_signals.py b/tests/test_signals.py new file mode 100644 index 0000000..e94b289 --- /dev/null +++ b/tests/test_signals.py @@ -0,0 +1,85 @@ +import time +import uuid +from unittest import mock + +import pytest +from django.db import DEFAULT_DB_ALIAS, connections +from django.tasks import TaskResult, TaskResultStatus +from django.tasks.backends.immediate import ImmediateBackend +from django.tasks.signals import task_finished, task_started +from django.utils import timezone + +from tests.testapp.tasks import echo +from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread + + +def _task_result(task, *args) -> TaskResult: + """Build a READY `TaskResult` without touching the backend.""" + return TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TaskResultStatus.READY, + enqueued_at=timezone.now(), + started_at=None, + finished_at=None, + last_attempted_at=None, + args=list(args), + kwargs={}, + backend="default", + errors=[], + worker_ids=[], + ) + + +@pytest.fixture +def stale_connection(): + """Open a database connection and mark it as expired.""" + connection = connections[DEFAULT_DB_ALIAS] + connection.ensure_connection() + connection.close_at = time.monotonic() - 1 + yield connection + connection.close_at = None + connection.close() + + +@pytest.mark.django_db(transaction=True) +class TestCloseTaskDatabaseConnection: + def test_task_started__closes_stale_connection(self, stale_connection): + """A stale connection is closed when a task starts.""" + with mock.patch.object( + stale_connection, "close", wraps=stale_connection.close + ) as spy: + task_started.send(TaskExecutor, task_result=_task_result(echo)) + + spy.assert_called_once() + + def test_task_finished__closes_stale_connection(self, stale_connection): + """A stale connection is closed when a task finished.""" + with mock.patch.object( + stale_connection, "close", wraps=stale_connection.close + ) as spy: + task_finished.send(TaskExecutor, task_result=_task_result(echo)) + + spy.assert_called_once() + + def test_task_started__ignores_other_senders(self, stale_connection): + """Signals from other executors do not close the connection.""" + with mock.patch.object(stale_connection, "close") as spy: + task_started.send(ImmediateBackend, task_result=_task_result(echo)) + + spy.assert_not_called() + + def test_execute_task_result__closes_stale_connection(self, stale_connection): + """Executing a task closes stale connections via lifecycle signals.""" + with mock.patch.object( + stale_connection, "close", wraps=stale_connection.close + ) as spy: + worker = WorkerProcess( + thread_count=1, backend_alias="default", queues=("default",) + ) + thread = WorkerThread(worker=worker, index=0, backend=None) + + result = thread.execute_task_result(_task_result(echo, 42)) + + assert result.status is TaskResultStatus.SUCCESSFUL + spy.assert_called() diff --git a/threadmill/apps.py b/threadmill/apps.py new file mode 100644 index 0000000..1b0a367 --- /dev/null +++ b/threadmill/apps.py @@ -0,0 +1,8 @@ +from django.apps import AppConfig + + +class ThreadmillConfig(AppConfig): + name = "threadmill" + + def ready(self) -> None: + from . import signals # noqa: F401 diff --git a/threadmill/signals.py b/threadmill/signals.py new file mode 100644 index 0000000..57bb710 --- /dev/null +++ b/threadmill/signals.py @@ -0,0 +1,10 @@ +from django.db import close_old_connections +from django.dispatch import receiver +from django.tasks.signals import task_finished, task_started + +from .executor import TaskExecutor + + +@receiver([task_started, task_finished], sender=TaskExecutor) +def close_task_database_connection(sender, task_result, **kwargs): + close_old_connections()