Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions tests/test_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
boom_retry_raises,
boom_retry_thrice,
boom_with_retry,
count_users,
echo,
)
from threadmill.backends.base import Broker
Expand Down Expand Up @@ -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])
Expand Down
2 changes: 1 addition & 1 deletion tests/testapp/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@
DATABASES = {
"default": {
"ENGINE": "django.db.backends.sqlite3",
"NAME": ":memory:",
"NAME": BASE_DIR / "db.sqlite3",
}
}

Expand Down
6 changes: 6 additions & 0 deletions tests/testapp/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
2 changes: 2 additions & 0 deletions threadmill/executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down