diff --git a/tests/test_executor.py b/tests/test_executor.py index 456e376..7009381 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -22,6 +22,7 @@ boom_retry_raises, boom_retry_thrice, boom_with_retry, + count_users, echo, ) from threadmill.backends.base import Broker @@ -158,6 +159,29 @@ def test_run__processes_enqueued_tasks_end_to_end(self): assert {r.id for r in results} == {r.id for r in enqueued} assert all(r.status == TaskResultStatus.SUCCESSFUL for r in results) + def test_run__executes_model_task_in_spawned_worker(self): + """run() executes a model-accessing task in a spawned worker process.""" + original_start_method = multiprocessing.get_start_method() + multiprocessing.set_start_method("spawn", force=True) + try: + enqueued = default_task_backend.enqueue(count_users) + executor = TaskExecutor( + backend=default_task_backend, + workers=1, + threads=1, + queues=("default",), + ) + run_thread = threading.Thread(target=executor.run, daemon=True) + run_thread.start() + time.sleep(3) + executor.shutdown() + run_thread.join(timeout=5) + assert not run_thread.is_alive() + result = default_task_backend.get_result(enqueued.id) + assert result.status == TaskResultStatus.SUCCESSFUL + finally: + multiprocessing.set_start_method(original_start_method, force=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]) diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index 1d4e2fd..dd4767c 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -82,7 +82,7 @@ DATABASES = { "default": { "ENGINE": "django.db.backends.sqlite3", - "NAME": ":memory:", + "NAME": BASE_DIR / "db.sqlite3", } } diff --git a/tests/testapp/tasks.py b/tests/testapp/tasks.py index 0be41ac..af578f6 100644 --- a/tests/testapp/tasks.py +++ b/tests/testapp/tasks.py @@ -22,6 +22,12 @@ def boom(): raise ValueError("boom") +@task() +def count_users(): + """Count all users in the database (tests model access in workers).""" + from django.contrib.auth.models import User # noqa + + @task(queue_name="compute") def compute_workload(): """Calculate the first 1000 prime numbers.""" diff --git a/threadmill/executor.py b/threadmill/executor.py index cfd4d64..f9490b3 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -15,6 +15,7 @@ from queue import Empty from traceback import format_exception +import django from django.tasks import TaskResult, task_backends from django.tasks.base import TaskContext, TaskError, TaskResultStatus from django.tasks.signals import task_finished, task_started @@ -146,6 +147,7 @@ def __init__( def run(self) -> None: """Start consumer execution inside this process.""" + django.setup() logger.info("Starting worker process %s", self.name) self.lock = threading.Lock() self.expired = threading.Event()