445 lines
14 KiB
Python
445 lines
14 KiB
Python
import datetime
|
|
import multiprocessing
|
|
import random
|
|
import signal
|
|
import time
|
|
from unittest import mock
|
|
|
|
from django.utils.timezone import now as tz_now
|
|
import pytest
|
|
|
|
from awx.main.models import Job, WorkflowJob, Instance
|
|
from awx.main.dispatch import reaper
|
|
from awx.main.dispatch.pool import StatefulPoolWorker, WorkerPool, AutoscalePool
|
|
from awx.main.dispatch.publish import task
|
|
from awx.main.dispatch.worker import BaseWorker, TaskWorker
|
|
|
|
|
|
'''
|
|
Prevent logger.<warn, debug, error> calls from triggering database operations
|
|
'''
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _disable_database_settings(mocker):
|
|
m = mocker.patch('awx.conf.settings.SettingsWrapper.all_supported_settings', new_callable=mock.PropertyMock)
|
|
m.return_value = []
|
|
|
|
|
|
def restricted(a, b):
|
|
raise AssertionError("This code should not run because it isn't decorated with @task")
|
|
|
|
|
|
@task()
|
|
def add(a, b):
|
|
return a + b
|
|
|
|
|
|
class BaseTask(object):
|
|
|
|
def add(self, a, b):
|
|
return add(a, b)
|
|
|
|
|
|
class Restricted(object):
|
|
def run(self, a, b):
|
|
raise AssertionError("This code should not run because it isn't decorated with @task")
|
|
|
|
|
|
@task()
|
|
class Adder(BaseTask):
|
|
def run(self, a, b):
|
|
return super(Adder, self).add(a, b)
|
|
|
|
|
|
@task(queue='hard-math')
|
|
def multiply(a, b):
|
|
return a * b
|
|
|
|
|
|
class SimpleWorker(BaseWorker):
|
|
|
|
def perform_work(self, body, *args):
|
|
pass
|
|
|
|
|
|
class ResultWriter(BaseWorker):
|
|
|
|
def perform_work(self, body, result_queue):
|
|
result_queue.put(body + '!!!')
|
|
|
|
|
|
class SlowResultWriter(BaseWorker):
|
|
|
|
def perform_work(self, body, result_queue):
|
|
time.sleep(3)
|
|
super(SlowResultWriter, self).perform_work(body, result_queue)
|
|
|
|
|
|
@pytest.mark.usefixtures("disable_database_settings")
|
|
class TestPoolWorker:
|
|
|
|
def setup_method(self, test_method):
|
|
self.worker = StatefulPoolWorker(1000, self.tick, tuple())
|
|
|
|
def tick(self):
|
|
self.worker.finished.put(self.worker.queue.get()['uuid'])
|
|
time.sleep(.5)
|
|
|
|
def test_qsize(self):
|
|
assert self.worker.qsize == 0
|
|
for i in range(3):
|
|
self.worker.put({'task': 'abc123'})
|
|
assert self.worker.qsize == 3
|
|
|
|
def test_put(self):
|
|
assert len(self.worker.managed_tasks) == 0
|
|
assert self.worker.messages_finished == 0
|
|
self.worker.put({'task': 'abc123'})
|
|
|
|
assert len(self.worker.managed_tasks) == 1
|
|
assert self.worker.messages_sent == 1
|
|
|
|
def test_managed_tasks(self):
|
|
self.worker.put({'task': 'abc123'})
|
|
self.worker.calculate_managed_tasks()
|
|
assert len(self.worker.managed_tasks) == 1
|
|
|
|
self.tick()
|
|
self.worker.calculate_managed_tasks()
|
|
assert len(self.worker.managed_tasks) == 0
|
|
|
|
def test_current_task(self):
|
|
self.worker.put({'task': 'abc123'})
|
|
assert self.worker.current_task['task'] == 'abc123'
|
|
|
|
def test_quit(self):
|
|
self.worker.quit()
|
|
assert self.worker.queue.get() == 'QUIT'
|
|
|
|
def test_idle_busy(self):
|
|
assert self.worker.idle is True
|
|
assert self.worker.busy is False
|
|
self.worker.put({'task': 'abc123'})
|
|
assert self.worker.busy is True
|
|
assert self.worker.idle is False
|
|
|
|
|
|
@pytest.mark.django_db
|
|
class TestWorkerPool:
|
|
|
|
def setup_method(self, test_method):
|
|
self.pool = WorkerPool(min_workers=3)
|
|
|
|
def teardown_method(self, test_method):
|
|
self.pool.stop(signal.SIGTERM)
|
|
|
|
def test_worker(self):
|
|
self.pool.init_workers(SimpleWorker().work_loop)
|
|
assert len(self.pool) == 3
|
|
for worker in self.pool.workers:
|
|
assert worker.messages_sent == 0
|
|
assert worker.alive is True
|
|
|
|
def test_single_task(self):
|
|
self.pool.init_workers(SimpleWorker().work_loop)
|
|
self.pool.write(0, 'xyz')
|
|
assert self.pool.workers[0].messages_sent == 1 # worker at index 0 handled one task
|
|
assert self.pool.workers[1].messages_sent == 0
|
|
assert self.pool.workers[2].messages_sent == 0
|
|
|
|
def test_queue_preference(self):
|
|
self.pool.init_workers(SimpleWorker().work_loop)
|
|
self.pool.write(2, 'xyz')
|
|
assert self.pool.workers[0].messages_sent == 0
|
|
assert self.pool.workers[1].messages_sent == 0
|
|
assert self.pool.workers[2].messages_sent == 1 # worker at index 2 handled one task
|
|
|
|
def test_worker_processing(self):
|
|
result_queue = multiprocessing.Queue()
|
|
self.pool.init_workers(ResultWriter().work_loop, result_queue)
|
|
for i in range(10):
|
|
self.pool.write(
|
|
random.choice(range(len(self.pool))),
|
|
'Hello, Worker {}'.format(i)
|
|
)
|
|
all_messages = [result_queue.get(timeout=1) for i in range(10)]
|
|
all_messages.sort()
|
|
assert all_messages == [
|
|
'Hello, Worker {}!!!'.format(i)
|
|
for i in range(10)
|
|
]
|
|
|
|
total_handled = sum([worker.messages_sent for worker in self.pool.workers])
|
|
assert total_handled == 10
|
|
|
|
|
|
@pytest.mark.django_db
|
|
class TestAutoScaling:
|
|
|
|
def setup_method(self, test_method):
|
|
self.pool = AutoscalePool(min_workers=2, max_workers=10)
|
|
|
|
def teardown_method(self, test_method):
|
|
self.pool.stop(signal.SIGTERM)
|
|
|
|
def test_scale_up(self):
|
|
result_queue = multiprocessing.Queue()
|
|
self.pool.init_workers(SlowResultWriter().work_loop, result_queue)
|
|
|
|
# start with two workers, write an event to each worker and make it busy
|
|
assert len(self.pool) == 2
|
|
for i, w in enumerate(self.pool.workers):
|
|
w.put('Hello, Worker {}'.format(0))
|
|
assert len(self.pool) == 2
|
|
|
|
# wait for the subprocesses to start working on their tasks and be marked busy
|
|
time.sleep(1)
|
|
assert self.pool.should_grow
|
|
|
|
# write a third message, expect a new worker to spawn because all
|
|
# workers are busy
|
|
self.pool.write(0, 'Hello, Worker {}'.format(2))
|
|
assert len(self.pool) == 3
|
|
|
|
def test_scale_down(self):
|
|
self.pool.init_workers(ResultWriter().work_loop, multiprocessing.Queue())
|
|
|
|
# start with two workers, and scale up to 10 workers
|
|
assert len(self.pool) == 2
|
|
for i in range(8):
|
|
self.pool.up()
|
|
assert len(self.pool) == 10
|
|
|
|
# cleanup should scale down to 8 workers
|
|
with mock.patch('awx.main.dispatch.reaper.reap') as reap:
|
|
self.pool.cleanup()
|
|
reap.assert_called()
|
|
assert len(self.pool) == 2
|
|
|
|
def test_max_scale_up(self):
|
|
self.pool.init_workers(ResultWriter().work_loop, multiprocessing.Queue())
|
|
|
|
assert len(self.pool) == 2
|
|
for i in range(25):
|
|
self.pool.up()
|
|
assert self.pool.max_workers == 10
|
|
assert self.pool.full is True
|
|
assert len(self.pool) == 10
|
|
|
|
def test_equal_worker_distribution(self):
|
|
# if all workers are busy, spawn new workers *before* adding messages
|
|
# to an existing queue
|
|
self.pool.init_workers(SlowResultWriter().work_loop, multiprocessing.Queue)
|
|
|
|
# start with two workers, write an event to each worker and make it busy
|
|
assert len(self.pool) == 2
|
|
for i in range(10):
|
|
self.pool.write(0, 'Hello, World!')
|
|
assert len(self.pool) == 10
|
|
for w in self.pool.workers:
|
|
assert w.busy
|
|
assert len(w.managed_tasks) == 1
|
|
|
|
# the queue is full at 10, the _next_ write should put the message into
|
|
# a worker's backlog
|
|
assert len(self.pool) == 10
|
|
for w in self.pool.workers:
|
|
assert w.messages_sent == 1
|
|
self.pool.write(0, 'Hello, World!')
|
|
assert len(self.pool) == 10
|
|
assert self.pool.workers[0].messages_sent == 2
|
|
|
|
def test_lost_worker_autoscale(self):
|
|
# if a worker exits, it should be replaced automatically up to min_workers
|
|
self.pool.init_workers(ResultWriter().work_loop, multiprocessing.Queue())
|
|
|
|
# start with two workers, kill one of them
|
|
assert len(self.pool) == 2
|
|
assert not self.pool.should_grow
|
|
alive_pid = self.pool.workers[1].pid
|
|
self.pool.workers[0].process.terminate()
|
|
time.sleep(1) # wait a moment for sigterm
|
|
|
|
# clean up and the dead worker
|
|
with mock.patch('awx.main.dispatch.reaper.reap') as reap:
|
|
self.pool.cleanup()
|
|
reap.assert_called()
|
|
assert len(self.pool) == 1
|
|
assert self.pool.workers[0].pid == alive_pid
|
|
|
|
# the next queue write should replace the lost worker
|
|
self.pool.write(0, 'Hello, Worker')
|
|
assert len(self.pool) == 2
|
|
|
|
|
|
@pytest.mark.usefixtures("disable_database_settings")
|
|
class TestTaskDispatcher:
|
|
|
|
@property
|
|
def tm(self):
|
|
return TaskWorker()
|
|
|
|
def test_function_dispatch(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'awx.main.tests.functional.test_dispatch.add',
|
|
'args': [2, 2]
|
|
})
|
|
assert result == 4
|
|
|
|
def test_function_dispatch_must_be_decorated(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'awx.main.tests.functional.test_dispatch.restricted',
|
|
'args': [2, 2]
|
|
})
|
|
assert isinstance(result, ValueError)
|
|
assert str(result) == 'awx.main.tests.functional.test_dispatch.restricted is not decorated with @task()' # noqa
|
|
|
|
def test_method_dispatch(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'awx.main.tests.functional.test_dispatch.Adder',
|
|
'args': [2, 2]
|
|
})
|
|
assert result == 4
|
|
|
|
def test_method_dispatch_must_be_decorated(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'awx.main.tests.functional.test_dispatch.Restricted',
|
|
'args': [2, 2]
|
|
})
|
|
assert isinstance(result, ValueError)
|
|
assert str(result) == 'awx.main.tests.functional.test_dispatch.Restricted is not decorated with @task()' # noqa
|
|
|
|
def test_python_function_cannot_be_imported(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'os.system',
|
|
'args': ['ls'],
|
|
})
|
|
assert isinstance(result, ValueError)
|
|
assert str(result) == 'os.system is not a valid awx task' # noqa
|
|
|
|
def test_undefined_function_cannot_be_imported(self):
|
|
result = self.tm.perform_work({
|
|
'task': 'awx.foo.bar'
|
|
})
|
|
assert isinstance(result, ModuleNotFoundError)
|
|
assert str(result) == "No module named 'awx.foo'" # noqa
|
|
|
|
|
|
class TestTaskPublisher:
|
|
|
|
def test_function_callable(self):
|
|
assert add(2, 2) == 4
|
|
|
|
def test_method_callable(self):
|
|
assert Adder().run(2, 2) == 4
|
|
|
|
def test_function_apply_async(self):
|
|
message, queue = add.apply_async([2, 2], queue='foobar')
|
|
assert message['args'] == [2, 2]
|
|
assert message['kwargs'] == {}
|
|
assert message['task'] == 'awx.main.tests.functional.test_dispatch.add'
|
|
assert queue == 'foobar'
|
|
|
|
def test_method_apply_async(self):
|
|
message, queue = Adder.apply_async([2, 2], queue='foobar')
|
|
assert message['args'] == [2, 2]
|
|
assert message['kwargs'] == {}
|
|
assert message['task'] == 'awx.main.tests.functional.test_dispatch.Adder'
|
|
assert queue == 'foobar'
|
|
|
|
def test_apply_async_queue_required(self):
|
|
with pytest.raises(ValueError) as e:
|
|
message, queue = add.apply_async([2, 2])
|
|
assert "awx.main.tests.functional.test_dispatch.add: Queue value required and may not be None" == e.value.args[0]
|
|
|
|
def test_queue_defined_in_task_decorator(self):
|
|
message, queue = multiply.apply_async([2, 2])
|
|
assert queue == 'hard-math'
|
|
|
|
def test_queue_overridden_from_task_decorator(self):
|
|
message, queue = multiply.apply_async([2, 2], queue='not-so-hard')
|
|
assert queue == 'not-so-hard'
|
|
|
|
def test_apply_with_callable_queuename(self):
|
|
message, queue = add.apply_async([2, 2], queue=lambda: 'called')
|
|
assert queue == 'called'
|
|
|
|
|
|
yesterday = tz_now() - datetime.timedelta(days=1)
|
|
|
|
|
|
@pytest.mark.django_db
|
|
class TestJobReaper(object):
|
|
|
|
@pytest.mark.parametrize('status, execution_node, controller_node, modified, fail', [
|
|
('running', '', '', None, False), # running, not assigned to the instance
|
|
('running', 'awx', '', None, True), # running, has the instance as its execution_node
|
|
('running', '', 'awx', None, True), # running, has the instance as its controller_node
|
|
('waiting', '', '', None, False), # waiting, not assigned to the instance
|
|
('waiting', 'awx', '', None, False), # waiting, was edited less than a minute ago
|
|
('waiting', '', 'awx', None, False), # waiting, was edited less than a minute ago
|
|
('waiting', 'awx', '', yesterday, True), # waiting, assigned to the execution_node, stale
|
|
('waiting', '', 'awx', yesterday, True), # waiting, assigned to the controller_node, stale
|
|
])
|
|
def test_should_reap(self, status, fail, execution_node, controller_node, modified):
|
|
i = Instance(hostname='awx')
|
|
i.save()
|
|
j = Job(
|
|
status=status,
|
|
execution_node=execution_node,
|
|
controller_node=controller_node,
|
|
start_args='SENSITIVE',
|
|
)
|
|
j.save()
|
|
if modified:
|
|
# we have to edit the modification time _without_ calling save()
|
|
# (because .save() overwrites it to _now_)
|
|
Job.objects.filter(id=j.id).update(modified=modified)
|
|
reaper.reap(i)
|
|
job = Job.objects.first()
|
|
if fail:
|
|
assert job.status == 'failed'
|
|
assert 'marked as failed' in job.job_explanation
|
|
assert job.start_args == ''
|
|
else:
|
|
assert job.status == status
|
|
|
|
@pytest.mark.parametrize('excluded_uuids, fail', [
|
|
(['abc123'], False),
|
|
([], True),
|
|
])
|
|
def test_do_not_reap_excluded_uuids(self, excluded_uuids, fail):
|
|
i = Instance(hostname='awx')
|
|
i.save()
|
|
j = Job(
|
|
status='running',
|
|
execution_node='awx',
|
|
controller_node='',
|
|
start_args='SENSITIVE',
|
|
celery_task_id='abc123',
|
|
)
|
|
j.save()
|
|
|
|
# if the UUID is excluded, don't reap it
|
|
reaper.reap(i, excluded_uuids=excluded_uuids)
|
|
job = Job.objects.first()
|
|
if fail:
|
|
assert job.status == 'failed'
|
|
assert 'marked as failed' in job.job_explanation
|
|
assert job.start_args == ''
|
|
else:
|
|
assert job.status == 'running'
|
|
|
|
def test_workflow_does_not_reap(self):
|
|
i = Instance(hostname='awx')
|
|
i.save()
|
|
j = WorkflowJob(
|
|
status='running',
|
|
execution_node='awx'
|
|
)
|
|
j.save()
|
|
reaper.reap(i)
|
|
|
|
assert WorkflowJob.objects.first().status == 'running'
|