Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def setUp(self):
AsyncHTTPTestCase.setUp(self)
self.is_worker = ControlHandler.is_worker
ControlHandler.is_worker = lambda *args: True
from tests.unit import AsyncHTTPTestCase
inspect_response = {
'celery@worker1': [
"tasks.add",
"tasks.sleep"
],
}
empty_inspect_response = {
'celery@worker1': []
}
@mock.patch.object(ControlHandler, 'INSPECT_METHODS',
new_callable=mock.PropertyMock,
return_value=['inspect_method'])
class ListWorkersTest(AsyncHTTPTestCase):
def test_refresh_cache(self, m_inspect):
celery = self._app.capp
celery.control.inspect = mock.Mock()
celery.control.inspect.return_value.inspect_method = mock.Mock(
return_value=inspect_response
)
r = self.get('/api/workers?refresh=1')
celery.control.inspect.assert_called_once_with(
timeout=1,
destination=None
)
taskname, soft, hard)
destination = [workername] if workername is not None else None
response = celery.control.time_limit(taskname, reply=True,
hard=hard, soft=soft,
destination=destination)
if response and 'ok' in response[0][workername]:
self.write(dict(message=response[0][workername]['ok']))
else:
logging.error(response)
self.set_status(403)
self.write("Failed to set timeouts: '%s'" %
self.error_reason(taskname, response))
class TaskRateLimit(ControlHandler):
@web.authenticated
def post(self, taskname):
"""
Change rate limit for a task
**Example request**:
.. sourcecode:: http
POST /api/task/rate-limit/tasks.sleep HTTP/1.1
Content-Length: 41
Content-Type: application/x-www-form-urlencoded; charset=utf-8
Host: localhost:5555
ratelimit=200&workername=celery%40worker1
"message": "Revoked '1480b55c-b8b2-462c-985e-24af3e9158f9'"
}
:query terminate: terminate the task if it is running
:reqheader Authorization: optional OAuth token to authenticate
:statuscode 200: no error
:statuscode 401: unauthorized request
"""
logging.info("Revoking task '%s'", taskid)
celery = self.application.celery_app
terminate = self.get_argument('terminate', default=False, type=bool)
celery.control.revoke(taskid, terminate=terminate)
self.write(dict(message="Revoked '%s'" % taskid))
class TaskTimout(ControlHandler):
@web.authenticated
def post(self, taskname):
"""
Change soft and hard time limits for a task
**Example request**:
.. sourcecode:: http
POST /api/task/timeout/tasks.sleep HTTP/1.1
Content-Length: 44
Content-Type: application/x-www-form-urlencoded; charset=utf-8
Host: localhost:5555
soft=30&hard=100&workername=celery%40worker1
response = self.capp.control.broadcast(
'autoscale', arguments={'min': min, 'max': max},
destination=[workername], reply=True)
if response and 'ok' in response[0][workername]:
self.write(dict(message="Autoscaling '%s' worker "
"(min=%s, max=%s)" % (
workername, min, max)))
else:
logger.error(response)
self.set_status(403)
self.write("Failed to autoscale '%s' worker: %s" % (
workername, self.error_reason(workername, response)
))
class WorkerQueueAddConsumer(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Start consuming from a queue
**Example request**:
.. sourcecode:: http
POST /api/worker/queue/add-consumer/celery@worker2?queue=sample-queue
Content-Length: 0
Content-Type: application/x-www-form-urlencoded; charset=utf-8
Host: localhost:5555
**Example response**:
logger.info("Shrinking '%s' worker's pool by '%s'", workername, n)
response = self.capp.control.pool_shrink(
n=n, reply=True, destination=[workername])
if response and 'ok' in response[0][workername]:
self.write(dict(message="Shrinking '%s' worker's pool by %s" % (
workername, n)))
else:
logger.error(response)
self.set_status(403)
self.write("Failed to shrink '%s' worker's pool: %s" % (
workername, self.error_reason(workername, response)
))
class WorkerPoolAutoscale(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Autoscale worker pool
**Example request**:
.. sourcecode:: http
POST /api/worker/pool/autoscale/celery@worker2?min=3&max=10 HTTP/1.1
Content-Length: 0
Content-Type: application/x-www-form-urlencoded; charset=utf-8
Host: localhost:5555
**Example response**:
response = celery.control.broadcast('pool_restart',
arguments={'reload': False},
destination=[workername],
reply=True)
if response and 'ok' in response[0][workername]:
self.write(dict(
message="Restarting '%s' worker's pool" % workername))
else:
logging.error(response)
self.set_status(403)
self.write("Failed to restart the '%s' pool: %s" % (
workername, self.error_reason(workername, response)
))
class WorkerPoolGrow(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Grow worker's pool
**Example request**:
.. sourcecode:: http
POST /api/worker/pool/grow/celery@worker2?n=3 HTTP/1.1
Content-Length: 0
Host: localhost:5555
**Example response**:
.. sourcecode:: http
class ControlHandler(BaseHandler):
def is_worker(self, name):
return WorkersModel.is_worker(self.application, name)
def error_reason(self, workername, response):
"extracts error message from response"
for r in response:
try:
return r[workername].get('error', 'Unknown error')
except KeyError:
pass
class WorkerShutDown(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Shut down a worker
**Example request**:
.. sourcecode:: http
POST /api/worker/shutdown/celery@worker2 HTTP/1.1
Content-Length: 0
Host: localhost:5555
**Example response**:
.. sourcecode:: http
arguments={'min': min, 'max': max},
destination=[workername],
reply=True)
if response and 'ok' in response[0][workername]:
self.write(dict(message="Autoscaling '%s' worker "
"(min=%s, max=%s)" % (
workername, min, max)))
else:
logging.error(response)
self.set_status(403)
self.write("Failed to autoscale '%s' worker: %s" % (
workername, self.error_reason(workername, response)
))
class WorkerQueueAddConsumer(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Start consuming from a queue
**Example request**:
.. sourcecode:: http
POST /api/worker/queue/add-consumer/celery@worker2?queue=sample-queue
Content-Length: 0
Content-Type: application/x-www-form-urlencoded; charset=utf-8
Host: localhost:5555
**Example response**:
n = self.get_argument('n', default=1, type=int)
logging.info("Growing '%s' worker's pool by '%s'", workername, n)
response = celery.control.pool_grow(n=n, reply=True,
destination=[workername])
if response and 'ok' in response[0][workername]:
self.write(dict(
message="Growing '%s' worker's pool by %s" % (workername, n)))
else:
logging.error(response)
self.set_status(403)
self.write("Failed to grow '%s' worker's pool" % (
workername, self.error_reason(workername, response)))
class WorkerPoolShrink(ControlHandler):
@web.authenticated
def post(self, workername):
"""
Shrink worker's pool
**Example request**:
.. sourcecode:: http
POST /api/worker/pool/shrink/celery@worker2 HTTP/1.1
Content-Length: 0
Host: localhost:5555
**Example response**:
.. sourcecode:: http