p1-taskqueue 0.1.8__tar.gz → 0.1.9__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of p1-taskqueue might be problematic. Click here for more details.
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/PKG-INFO +1 -1
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/pyproject.toml +1 -1
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/p1_taskqueue.egg-info/PKG-INFO +1 -1
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/p1_taskqueue.egg-info/SOURCES.txt +1 -0
- p1_taskqueue-0.1.9/src/taskqueue/libs/helper_test.py +178 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/tests/test_cmanager.py +45 -43
- p1_taskqueue-0.1.9/tests/test_helper_test_functions.py +350 -0
- p1_taskqueue-0.1.8/src/taskqueue/libs/helper_test.py +0 -93
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/README.md +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/setup.cfg +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/p1_taskqueue.egg-info/dependency_links.txt +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/p1_taskqueue.egg-info/requires.txt +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/p1_taskqueue.egg-info/top_level.txt +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/taskqueue/__init__.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/taskqueue/celery_app.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/taskqueue/cmanager.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/src/taskqueue/libs/__init__.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/tests/test_celery_app.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/tests/test_return_values.py +0 -0
- {p1_taskqueue-0.1.8 → p1_taskqueue-0.1.9}/tests/test_test_utils.py +0 -0
|
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "p1-taskqueue"
|
|
7
7
|
# DO NOT CHANGE THIS VERSION - it gets automatically replaced by CI/CD with the git tag version
|
|
8
|
-
version = "0.1.
|
|
8
|
+
version = "0.1.9"
|
|
9
9
|
description = "A Task Queue Wrapper for Dekoruma Backend"
|
|
10
10
|
authors = [
|
|
11
11
|
{name = "Chalvin", email = "engineering@dekoruma.com"}
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
from typing import List
|
|
3
|
+
|
|
4
|
+
from celery import current_app
|
|
5
|
+
from kombu.serialization import loads
|
|
6
|
+
|
|
7
|
+
logger = logging.getLogger(__name__)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def clear_all_celery_queues():
|
|
11
|
+
app = current_app
|
|
12
|
+
all_queue_names = list(app.amqp.queues.keys())
|
|
13
|
+
with app.connection_for_read() as conn:
|
|
14
|
+
with conn.channel() as chan:
|
|
15
|
+
for queue_name in all_queue_names:
|
|
16
|
+
queue = app.amqp.queues[queue_name](chan)
|
|
17
|
+
queue.purge()
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def celery_worker_burst(include_func_names: List[str], channel: str = "default"):
|
|
21
|
+
# This doesn't use celery as celery doesn't support filtering out functions
|
|
22
|
+
# this use kombu to get the message from the queue and then execute the task manually
|
|
23
|
+
app = current_app
|
|
24
|
+
included_set = set(include_func_names)
|
|
25
|
+
processed_count = 0
|
|
26
|
+
executed_count = 0
|
|
27
|
+
|
|
28
|
+
try:
|
|
29
|
+
with app.connection_for_read() as conn:
|
|
30
|
+
with conn.channel() as chan:
|
|
31
|
+
queue = app.amqp.queues[channel](chan)
|
|
32
|
+
|
|
33
|
+
while True:
|
|
34
|
+
message = queue.get(no_ack=False)
|
|
35
|
+
if not message:
|
|
36
|
+
break
|
|
37
|
+
|
|
38
|
+
processed_count += 1
|
|
39
|
+
task_name = message.headers.get("task")
|
|
40
|
+
|
|
41
|
+
if not task_name or task_name not in app.tasks:
|
|
42
|
+
# task is not registered in celery
|
|
43
|
+
logger.warning(
|
|
44
|
+
f"Invalid task '{task_name}'. Skipping.")
|
|
45
|
+
message.ack()
|
|
46
|
+
continue
|
|
47
|
+
|
|
48
|
+
try:
|
|
49
|
+
task_obj = app.tasks[task_name]
|
|
50
|
+
accept = {"application/json",
|
|
51
|
+
"application/x-python-serialize"}
|
|
52
|
+
decoded_body = loads(
|
|
53
|
+
message.body, message.content_type, message.content_encoding, accept=accept
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
task_args = decoded_body[0] if decoded_body else []
|
|
57
|
+
task_kwargs = decoded_body[1] if len(
|
|
58
|
+
decoded_body) > 1 else {}
|
|
59
|
+
|
|
60
|
+
full_func_name = ""
|
|
61
|
+
if task_name.endswith("dynamic_function_executor"):
|
|
62
|
+
module_path = task_kwargs.get('module_path', '')
|
|
63
|
+
function_name = task_kwargs.get(
|
|
64
|
+
'function_name', '')
|
|
65
|
+
if module_path and function_name:
|
|
66
|
+
full_func_name = f"{module_path}.{function_name}"
|
|
67
|
+
elif task_name.endswith("dynamic_class_method_executor"):
|
|
68
|
+
module_path = task_kwargs.get('module_path', '')
|
|
69
|
+
class_name = task_kwargs.get('class_name', '')
|
|
70
|
+
method_name = task_kwargs.get('method_name', '')
|
|
71
|
+
if module_path and class_name and method_name:
|
|
72
|
+
full_func_name = f"{module_path}.{class_name}.{method_name}"
|
|
73
|
+
|
|
74
|
+
should_execute = full_func_name in included_set if full_func_name else False
|
|
75
|
+
|
|
76
|
+
if should_execute:
|
|
77
|
+
logger.info(f"Executing task: {full_func_name}")
|
|
78
|
+
message.ack()
|
|
79
|
+
task_obj.apply(args=task_args, kwargs=task_kwargs)
|
|
80
|
+
executed_count += 1
|
|
81
|
+
logger.info(
|
|
82
|
+
f"Successfully executed task: {full_func_name}")
|
|
83
|
+
else:
|
|
84
|
+
logger.info(
|
|
85
|
+
f"Skipping: {full_func_name or task_name}")
|
|
86
|
+
message.ack()
|
|
87
|
+
|
|
88
|
+
except Exception as e:
|
|
89
|
+
logger.error(
|
|
90
|
+
f"Failed to process task {task_name}: {type(e).__name__}: {e}")
|
|
91
|
+
if message and not message.acknowledged:
|
|
92
|
+
message.ack()
|
|
93
|
+
|
|
94
|
+
except Exception as e:
|
|
95
|
+
logger.error(
|
|
96
|
+
f"Failed to connect to queue {channel}: {type(e).__name__}: {e}")
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def get_queued_tasks(channel: str = "default"):
|
|
100
|
+
app = current_app
|
|
101
|
+
queued_tasks = []
|
|
102
|
+
|
|
103
|
+
try:
|
|
104
|
+
with app.connection_for_read() as conn:
|
|
105
|
+
with conn.channel() as chan:
|
|
106
|
+
queue = app.amqp.queues[channel](chan)
|
|
107
|
+
|
|
108
|
+
while True:
|
|
109
|
+
message = queue.get(no_ack=True)
|
|
110
|
+
if not message:
|
|
111
|
+
break
|
|
112
|
+
|
|
113
|
+
task_name = message.headers.get("task")
|
|
114
|
+
|
|
115
|
+
try:
|
|
116
|
+
accept = {"application/json",
|
|
117
|
+
"application/x-python-serialize"}
|
|
118
|
+
decoded_body = loads(
|
|
119
|
+
message.body, message.content_type, message.content_encoding, accept=accept
|
|
120
|
+
)
|
|
121
|
+
|
|
122
|
+
task_args = decoded_body[0] if decoded_body else []
|
|
123
|
+
task_kwargs = decoded_body[1] if len(
|
|
124
|
+
decoded_body) > 1 else {}
|
|
125
|
+
|
|
126
|
+
full_func_name = ""
|
|
127
|
+
if task_name and task_name.endswith("dynamic_function_executor"):
|
|
128
|
+
module_path = task_kwargs.get('module_path', '')
|
|
129
|
+
function_name = task_kwargs.get(
|
|
130
|
+
'function_name', '')
|
|
131
|
+
if module_path and function_name:
|
|
132
|
+
full_func_name = f"{module_path}.{function_name}"
|
|
133
|
+
elif task_name and task_name.endswith("dynamic_class_method_executor"):
|
|
134
|
+
module_path = task_kwargs.get('module_path', '')
|
|
135
|
+
class_name = task_kwargs.get('class_name', '')
|
|
136
|
+
method_name = task_kwargs.get('method_name', '')
|
|
137
|
+
if module_path and class_name and method_name:
|
|
138
|
+
full_func_name = f"{module_path}.{class_name}.{method_name}"
|
|
139
|
+
|
|
140
|
+
queued_tasks.append({
|
|
141
|
+
'task_name': task_name,
|
|
142
|
+
'full_func_name': full_func_name,
|
|
143
|
+
'args': task_args,
|
|
144
|
+
'kwargs': task_kwargs,
|
|
145
|
+
'headers': message.headers
|
|
146
|
+
})
|
|
147
|
+
|
|
148
|
+
except Exception as e:
|
|
149
|
+
logger.warning(f"Failed to decode message: {e}")
|
|
150
|
+
continue
|
|
151
|
+
|
|
152
|
+
except Exception as e:
|
|
153
|
+
logger.error(
|
|
154
|
+
f"Failed to get queued tasks from queue {channel}: {type(e).__name__}: {e}")
|
|
155
|
+
|
|
156
|
+
return queued_tasks
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def is_task_in_queue(expected_func_name: str, channel: str = "default"):
|
|
160
|
+
queued_tasks = get_queued_tasks(channel)
|
|
161
|
+
for task in queued_tasks:
|
|
162
|
+
if task['full_func_name'] == expected_func_name:
|
|
163
|
+
return True
|
|
164
|
+
logger.info(
|
|
165
|
+
f"Task {expected_func_name} not found in queue {channel}. Queued tasks: {[task['full_func_name'] for task in queued_tasks]}")
|
|
166
|
+
return False
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def assert_task_in_queue(expected_func_name: str, channel: str = "default", msg: str = None):
|
|
170
|
+
if not is_task_in_queue(expected_func_name, channel):
|
|
171
|
+
error_msg = msg or f"Task '{expected_func_name}' not found in queue '{channel}'"
|
|
172
|
+
raise AssertionError(error_msg)
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def assert_task_not_in_queue(expected_func_name: str, channel: str = "default", msg: str = None):
|
|
176
|
+
if is_task_in_queue(expected_func_name, channel):
|
|
177
|
+
error_msg = msg or f"Task '{expected_func_name}' found in queue '{channel}'"
|
|
178
|
+
raise AssertionError(error_msg)
|
|
@@ -13,7 +13,7 @@ from taskqueue.cmanager import K_MAX_RETRY_COUNT
|
|
|
13
13
|
from taskqueue.cmanager import taskqueue_class
|
|
14
14
|
|
|
15
15
|
|
|
16
|
-
class
|
|
16
|
+
class SampleClass:
|
|
17
17
|
|
|
18
18
|
def test_method(self):
|
|
19
19
|
pass
|
|
@@ -28,7 +28,7 @@ class TestClass:
|
|
|
28
28
|
|
|
29
29
|
|
|
30
30
|
@taskqueue_class
|
|
31
|
-
class
|
|
31
|
+
class SampleClassWithInit:
|
|
32
32
|
|
|
33
33
|
def __init__(self, name, age=0, **kwargs):
|
|
34
34
|
self.name = name
|
|
@@ -43,7 +43,7 @@ class TestClassWithInit:
|
|
|
43
43
|
|
|
44
44
|
|
|
45
45
|
@taskqueue_class
|
|
46
|
-
class
|
|
46
|
+
class SampleClassWithVarArgs:
|
|
47
47
|
|
|
48
48
|
def __init__(self, name, *args, **kwargs):
|
|
49
49
|
self.name = name
|
|
@@ -55,7 +55,7 @@ class TestClassWithVarArgs:
|
|
|
55
55
|
|
|
56
56
|
|
|
57
57
|
@taskqueue_class
|
|
58
|
-
class
|
|
58
|
+
class SampleClassWithComplexInit:
|
|
59
59
|
|
|
60
60
|
def __init__(self, required, optional=10, *extra, **options):
|
|
61
61
|
self.required = required
|
|
@@ -68,7 +68,7 @@ class TestClassWithComplexInit:
|
|
|
68
68
|
|
|
69
69
|
|
|
70
70
|
@taskqueue_class
|
|
71
|
-
class
|
|
71
|
+
class SampleClassWithDifferentParamNames:
|
|
72
72
|
|
|
73
73
|
def __init__(self, cognito_form_reimbursement_dict):
|
|
74
74
|
self.data = cognito_form_reimbursement_dict
|
|
@@ -88,20 +88,20 @@ class TestIsClassMethod:
|
|
|
88
88
|
assert result is False
|
|
89
89
|
|
|
90
90
|
def test__is_class_method_given_class_method_expect_return_true(self):
|
|
91
|
-
instance =
|
|
91
|
+
instance = SampleClass()
|
|
92
92
|
result = _is_class_method(instance.test_method)
|
|
93
93
|
assert result is True
|
|
94
94
|
|
|
95
95
|
def test__is_class_method_given_classmethod_decorator_expect_return_true(self):
|
|
96
|
-
result = _is_class_method(
|
|
96
|
+
result = _is_class_method(SampleClass.class_method)
|
|
97
97
|
assert result is True
|
|
98
98
|
|
|
99
99
|
def test__is_class_method_given_staticmethod_decorator_expect_return_false(self):
|
|
100
|
-
result = _is_class_method(
|
|
100
|
+
result = _is_class_method(SampleClass.static_method)
|
|
101
101
|
assert result is False
|
|
102
102
|
|
|
103
103
|
def test__is_class_method_given_unbound_method_expect_return_false(self):
|
|
104
|
-
result = _is_class_method(
|
|
104
|
+
result = _is_class_method(SampleClass.test_method)
|
|
105
105
|
assert result is False
|
|
106
106
|
|
|
107
107
|
|
|
@@ -176,7 +176,7 @@ class TestBuildDynamicTaskCall:
|
|
|
176
176
|
}
|
|
177
177
|
|
|
178
178
|
def test__build_dynamic_task_call_given_bound_method_expect_class_method_executor_task(self):
|
|
179
|
-
instance =
|
|
179
|
+
instance = SampleClass()
|
|
180
180
|
task_name, task_args, task_kwargs = _build_dynamic_task_call(
|
|
181
181
|
instance.test_method, 1, 2, key='value'
|
|
182
182
|
)
|
|
@@ -185,7 +185,7 @@ class TestBuildDynamicTaskCall:
|
|
|
185
185
|
assert task_args == []
|
|
186
186
|
assert task_kwargs == {
|
|
187
187
|
'module_path': 'tests.test_cmanager',
|
|
188
|
-
'class_name': '
|
|
188
|
+
'class_name': 'SampleClass',
|
|
189
189
|
'method_name': 'test_method',
|
|
190
190
|
'args': [1, 2],
|
|
191
191
|
'kwargs': {'key': 'value'},
|
|
@@ -373,48 +373,50 @@ class TestExtractInitArgs:
|
|
|
373
373
|
|
|
374
374
|
def test_extract_init_args_from_instance_given_instance_with_init_args_expect_args_extracted(self):
|
|
375
375
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
376
|
-
instance =
|
|
377
|
-
|
|
376
|
+
instance = SampleClassWithInit("John", age=25)
|
|
377
|
+
|
|
378
378
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
379
379
|
assert init_args == ["John"]
|
|
380
380
|
assert init_kwargs == {"age": 25}
|
|
381
381
|
|
|
382
382
|
def test_extract_init_args_from_instance_given_kwargs_only_expect_extracted(self):
|
|
383
383
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
384
|
-
instance =
|
|
385
|
-
|
|
384
|
+
instance = SampleClassWithInit("Jane", age=30)
|
|
385
|
+
|
|
386
386
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
387
387
|
assert init_args == ["Jane"]
|
|
388
388
|
assert init_kwargs == {"age": 30}
|
|
389
389
|
|
|
390
390
|
def test_extract_init_args_from_instance_given_no_init_expect_empty(self):
|
|
391
391
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
392
|
-
instance =
|
|
393
|
-
|
|
392
|
+
instance = SampleClass()
|
|
393
|
+
|
|
394
394
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
395
395
|
assert init_args == []
|
|
396
396
|
assert init_kwargs == {}
|
|
397
397
|
|
|
398
398
|
def test_extract_init_args_from_instance_given_var_args_expect_captured(self):
|
|
399
399
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
400
|
-
instance =
|
|
401
|
-
|
|
400
|
+
instance = SampleClassWithVarArgs(
|
|
401
|
+
"Alice", "extra1", "extra2", city="NYC")
|
|
402
|
+
|
|
402
403
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
403
404
|
assert init_args == ["Alice", "extra1", "extra2"]
|
|
404
405
|
assert init_kwargs == {"city": "NYC"}
|
|
405
406
|
|
|
406
407
|
def test_extract_init_args_from_instance_given_var_kwargs_expect_captured(self):
|
|
407
408
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
408
|
-
instance =
|
|
409
|
-
|
|
409
|
+
instance = SampleClassWithInit("Bob", age=30, city="LA", country="USA")
|
|
410
|
+
|
|
410
411
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
411
412
|
assert init_args == ["Bob"]
|
|
412
413
|
assert init_kwargs == {"age": 30, "city": "LA", "country": "USA"}
|
|
413
414
|
|
|
414
415
|
def test_extract_init_args_from_instance_given_complex_init_expect_all_captured(self):
|
|
415
416
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
416
|
-
instance =
|
|
417
|
-
|
|
417
|
+
instance = SampleClassWithComplexInit(
|
|
418
|
+
5, 15, 20, 25, multiplier=2, debug=True)
|
|
419
|
+
|
|
418
420
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
419
421
|
assert init_args == [5, 15, 20, 25]
|
|
420
422
|
assert init_kwargs == {"multiplier": 2, "debug": True}
|
|
@@ -422,10 +424,10 @@ class TestExtractInitArgs:
|
|
|
422
424
|
def test_extract_init_args_from_instance_given_different_param_names_expect_works(self):
|
|
423
425
|
from taskqueue.cmanager import _extract_init_args_from_instance
|
|
424
426
|
from collections import OrderedDict
|
|
425
|
-
|
|
427
|
+
|
|
426
428
|
test_dict = OrderedDict([("key1", "value1"), ("key2", "value2")])
|
|
427
|
-
instance =
|
|
428
|
-
|
|
429
|
+
instance = SampleClassWithDifferentParamNames(test_dict)
|
|
430
|
+
|
|
429
431
|
init_args, init_kwargs = _extract_init_args_from_instance(instance)
|
|
430
432
|
assert init_args == [test_dict]
|
|
431
433
|
assert init_kwargs == {}
|
|
@@ -434,8 +436,8 @@ class TestExtractInitArgs:
|
|
|
434
436
|
class TestTaskqueueClassDecorator:
|
|
435
437
|
|
|
436
438
|
def test_taskqueue_class_decorator_given_class_expect_init_args_captured(self):
|
|
437
|
-
instance =
|
|
438
|
-
|
|
439
|
+
instance = SampleClassWithInit("John", age=25)
|
|
440
|
+
|
|
439
441
|
assert hasattr(instance, '_taskqueue_init_args')
|
|
440
442
|
assert hasattr(instance, '_taskqueue_init_kwargs')
|
|
441
443
|
assert instance._taskqueue_init_args == ["John"]
|
|
@@ -444,8 +446,8 @@ class TestTaskqueueClassDecorator:
|
|
|
444
446
|
def test_taskqueue_class_decorator_given_different_param_names_expect_captured(self):
|
|
445
447
|
from collections import OrderedDict
|
|
446
448
|
test_dict = OrderedDict([("key1", "value1")])
|
|
447
|
-
instance =
|
|
448
|
-
|
|
449
|
+
instance = SampleClassWithDifferentParamNames(test_dict)
|
|
450
|
+
|
|
449
451
|
assert instance._taskqueue_init_args == [test_dict]
|
|
450
452
|
assert instance._taskqueue_init_kwargs == {}
|
|
451
453
|
assert instance.data == test_dict
|
|
@@ -454,7 +456,7 @@ class TestTaskqueueClassDecorator:
|
|
|
454
456
|
class TestBuildDynamicTaskCallWithInitArgs:
|
|
455
457
|
|
|
456
458
|
def test__build_dynamic_task_call_given_instance_with_init_args_expect_init_args_passed(self):
|
|
457
|
-
instance =
|
|
459
|
+
instance = SampleClassWithInit("John", age=25)
|
|
458
460
|
task_name, task_args, task_kwargs = _build_dynamic_task_call(
|
|
459
461
|
instance.process, "arg1", key='value'
|
|
460
462
|
)
|
|
@@ -463,7 +465,7 @@ class TestBuildDynamicTaskCallWithInitArgs:
|
|
|
463
465
|
assert task_args == []
|
|
464
466
|
assert task_kwargs == {
|
|
465
467
|
'module_path': 'tests.test_cmanager',
|
|
466
|
-
'class_name': '
|
|
468
|
+
'class_name': 'SampleClassWithInit',
|
|
467
469
|
'method_name': 'process',
|
|
468
470
|
'args': ['arg1'],
|
|
469
471
|
'kwargs': {'key': 'value'},
|
|
@@ -472,7 +474,7 @@ class TestBuildDynamicTaskCallWithInitArgs:
|
|
|
472
474
|
}
|
|
473
475
|
|
|
474
476
|
def test__build_dynamic_task_call_given_instance_without_decorator_expect_empty_defaults(self):
|
|
475
|
-
instance =
|
|
477
|
+
instance = SampleClass()
|
|
476
478
|
task_name, task_args, task_kwargs = _build_dynamic_task_call(
|
|
477
479
|
instance.test_method, 1, 2, key='value'
|
|
478
480
|
)
|
|
@@ -481,7 +483,7 @@ class TestBuildDynamicTaskCallWithInitArgs:
|
|
|
481
483
|
assert task_args == []
|
|
482
484
|
assert task_kwargs == {
|
|
483
485
|
'module_path': 'tests.test_cmanager',
|
|
484
|
-
'class_name': '
|
|
486
|
+
'class_name': 'SampleClass',
|
|
485
487
|
'method_name': 'test_method',
|
|
486
488
|
'args': [1, 2],
|
|
487
489
|
'kwargs': {'key': 'value'},
|
|
@@ -495,10 +497,10 @@ class TestDynamicClassMethodExecutorWithInitArgs:
|
|
|
495
497
|
|
|
496
498
|
def test_dynamic_class_method_executor_given_init_args_expect_instance_created_with_args(self):
|
|
497
499
|
from taskqueue.cmanager import dynamic_class_method_executor
|
|
498
|
-
|
|
500
|
+
|
|
499
501
|
result = dynamic_class_method_executor(
|
|
500
502
|
module_path="tests.test_cmanager",
|
|
501
|
-
class_name="
|
|
503
|
+
class_name="SampleClassWithInit",
|
|
502
504
|
method_name="process",
|
|
503
505
|
args=[],
|
|
504
506
|
kwargs={},
|
|
@@ -506,15 +508,15 @@ class TestDynamicClassMethodExecutorWithInitArgs:
|
|
|
506
508
|
init_kwargs={"age": 35},
|
|
507
509
|
retry=None
|
|
508
510
|
)
|
|
509
|
-
|
|
511
|
+
|
|
510
512
|
assert result is None
|
|
511
513
|
|
|
512
514
|
def test_dynamic_class_method_executor_given_no_init_args_expect_backward_compatible(self):
|
|
513
515
|
from taskqueue.cmanager import dynamic_class_method_executor
|
|
514
|
-
|
|
516
|
+
|
|
515
517
|
result = dynamic_class_method_executor(
|
|
516
518
|
module_path="tests.test_cmanager",
|
|
517
|
-
class_name="
|
|
519
|
+
class_name="SampleClass",
|
|
518
520
|
method_name="test_method",
|
|
519
521
|
args=[],
|
|
520
522
|
kwargs={},
|
|
@@ -522,15 +524,15 @@ class TestDynamicClassMethodExecutorWithInitArgs:
|
|
|
522
524
|
init_kwargs=None,
|
|
523
525
|
retry=None
|
|
524
526
|
)
|
|
525
|
-
|
|
527
|
+
|
|
526
528
|
assert result is None
|
|
527
529
|
|
|
528
530
|
def test_dynamic_class_method_executor_given_var_args_and_kwargs_expect_works(self):
|
|
529
531
|
from taskqueue.cmanager import dynamic_class_method_executor
|
|
530
|
-
|
|
532
|
+
|
|
531
533
|
result = dynamic_class_method_executor(
|
|
532
534
|
module_path="tests.test_cmanager",
|
|
533
|
-
class_name="
|
|
535
|
+
class_name="SampleClassWithVarArgs",
|
|
534
536
|
method_name="process",
|
|
535
537
|
args=[],
|
|
536
538
|
kwargs={},
|
|
@@ -538,5 +540,5 @@ class TestDynamicClassMethodExecutorWithInitArgs:
|
|
|
538
540
|
init_kwargs={"city": "NYC", "country": "USA"},
|
|
539
541
|
retry=None
|
|
540
542
|
)
|
|
541
|
-
|
|
543
|
+
|
|
542
544
|
assert result is None
|
|
@@ -0,0 +1,350 @@
|
|
|
1
|
+
"""Tests for the helper_test module."""
|
|
2
|
+
from unittest.mock import MagicMock
|
|
3
|
+
from unittest.mock import patch
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
from taskqueue.libs.helper_test import assert_task_in_queue
|
|
7
|
+
from taskqueue.libs.helper_test import assert_task_not_in_queue
|
|
8
|
+
from taskqueue.libs.helper_test import celery_worker_burst
|
|
9
|
+
from taskqueue.libs.helper_test import clear_all_celery_queues
|
|
10
|
+
from taskqueue.libs.helper_test import get_queued_tasks
|
|
11
|
+
from taskqueue.libs.helper_test import is_task_in_queue
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def setup_celery_connection_mocks(mock_current_app):
|
|
15
|
+
"""Helper to setup Celery connection and channel mocks."""
|
|
16
|
+
mock_conn = MagicMock()
|
|
17
|
+
mock_chan = MagicMock()
|
|
18
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
19
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
20
|
+
return mock_conn, mock_chan
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class TestHelperTest:
|
|
24
|
+
|
|
25
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
26
|
+
def test_clear_all_celery_queues_given_multiple_queues_expect_all_purged(self, mock_current_app):
|
|
27
|
+
mock_conn, mock_chan = setup_celery_connection_mocks(mock_current_app)
|
|
28
|
+
|
|
29
|
+
mock_queue1 = MagicMock()
|
|
30
|
+
mock_queue2 = MagicMock()
|
|
31
|
+
mock_queue_factory1 = MagicMock(return_value=mock_queue1)
|
|
32
|
+
mock_queue_factory2 = MagicMock(return_value=mock_queue2)
|
|
33
|
+
|
|
34
|
+
mock_current_app.amqp.queues.keys.return_value = ['queue1', 'queue2']
|
|
35
|
+
mock_current_app.amqp.queues = {
|
|
36
|
+
'queue1': mock_queue_factory1,
|
|
37
|
+
'queue2': mock_queue_factory2
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
clear_all_celery_queues()
|
|
41
|
+
|
|
42
|
+
mock_queue_factory1.assert_called_once_with(mock_chan)
|
|
43
|
+
mock_queue_factory2.assert_called_once_with(mock_chan)
|
|
44
|
+
mock_queue1.purge.assert_called_once()
|
|
45
|
+
mock_queue2.purge.assert_called_once()
|
|
46
|
+
|
|
47
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
48
|
+
def test_clear_all_celery_queues_given_empty_queues_expect_no_purge_calls(self, mock_current_app):
|
|
49
|
+
mock_conn, mock_chan = setup_celery_connection_mocks(mock_current_app)
|
|
50
|
+
|
|
51
|
+
mock_current_app.amqp.queues.keys.return_value = []
|
|
52
|
+
mock_current_app.amqp.queues = {}
|
|
53
|
+
|
|
54
|
+
clear_all_celery_queues()
|
|
55
|
+
|
|
56
|
+
mock_conn.channel.assert_called_once()
|
|
57
|
+
|
|
58
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
59
|
+
def test_get_queued_tasks_given_empty_queue_expect_empty_list(self, mock_current_app):
|
|
60
|
+
mock_queue = MagicMock()
|
|
61
|
+
mock_queue.get.return_value = None
|
|
62
|
+
mock_current_app.amqp.queues = {
|
|
63
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
64
|
+
|
|
65
|
+
result = get_queued_tasks('default')
|
|
66
|
+
|
|
67
|
+
assert result == []
|
|
68
|
+
|
|
69
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
70
|
+
@patch('taskqueue.libs.helper_test.loads')
|
|
71
|
+
def test_get_queued_tasks_given_function_executor_expect_correct_parsing(self, mock_loads, mock_current_app):
|
|
72
|
+
mock_message = MagicMock()
|
|
73
|
+
mock_message.headers = {
|
|
74
|
+
'task': 'taskqueue.cmanager.dynamic_function_executor'}
|
|
75
|
+
mock_message.body = b'mock_body'
|
|
76
|
+
mock_message.content_type = 'application/json'
|
|
77
|
+
mock_message.content_encoding = 'utf-8'
|
|
78
|
+
|
|
79
|
+
mock_loads.return_value = [[], {
|
|
80
|
+
'module_path': 'my.module',
|
|
81
|
+
'function_name': 'my_function',
|
|
82
|
+
'args': [],
|
|
83
|
+
'kwargs': {}
|
|
84
|
+
}]
|
|
85
|
+
|
|
86
|
+
mock_queue = MagicMock()
|
|
87
|
+
mock_queue.get.side_effect = [mock_message, None]
|
|
88
|
+
mock_current_app.amqp.queues = {
|
|
89
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
90
|
+
|
|
91
|
+
mock_conn = MagicMock()
|
|
92
|
+
mock_chan = MagicMock()
|
|
93
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
94
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
95
|
+
|
|
96
|
+
result = get_queued_tasks('default')
|
|
97
|
+
|
|
98
|
+
assert len(result) == 1
|
|
99
|
+
assert result[0]['task_name'] == 'taskqueue.cmanager.dynamic_function_executor'
|
|
100
|
+
assert result[0]['full_func_name'] == 'my.module.my_function'
|
|
101
|
+
assert result[0]['args'] == []
|
|
102
|
+
assert result[0]['kwargs']['module_path'] == 'my.module'
|
|
103
|
+
assert result[0]['kwargs']['function_name'] == 'my_function'
|
|
104
|
+
|
|
105
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
106
|
+
@patch('taskqueue.libs.helper_test.loads')
|
|
107
|
+
def test_get_queued_tasks_given_class_method_executor_expect_correct_parsing(self, mock_loads, mock_current_app):
|
|
108
|
+
mock_message = MagicMock()
|
|
109
|
+
mock_message.headers = {
|
|
110
|
+
'task': 'taskqueue.cmanager.dynamic_class_method_executor'}
|
|
111
|
+
mock_message.body = b'mock_body'
|
|
112
|
+
mock_message.content_type = 'application/json'
|
|
113
|
+
mock_message.content_encoding = 'utf-8'
|
|
114
|
+
|
|
115
|
+
mock_loads.return_value = [[], {
|
|
116
|
+
'module_path': 'my.module',
|
|
117
|
+
'class_name': 'MyClass',
|
|
118
|
+
'method_name': 'my_method',
|
|
119
|
+
'args': [],
|
|
120
|
+
'kwargs': {}
|
|
121
|
+
}]
|
|
122
|
+
|
|
123
|
+
mock_queue = MagicMock()
|
|
124
|
+
mock_queue.get.side_effect = [mock_message, None]
|
|
125
|
+
mock_current_app.amqp.queues = {
|
|
126
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
127
|
+
|
|
128
|
+
mock_conn = MagicMock()
|
|
129
|
+
mock_chan = MagicMock()
|
|
130
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
131
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
132
|
+
|
|
133
|
+
result = get_queued_tasks('default')
|
|
134
|
+
|
|
135
|
+
assert len(result) == 1
|
|
136
|
+
assert result[0]['task_name'] == 'taskqueue.cmanager.dynamic_class_method_executor'
|
|
137
|
+
assert result[0]['full_func_name'] == 'my.module.MyClass.my_method'
|
|
138
|
+
assert result[0]['args'] == []
|
|
139
|
+
assert result[0]['kwargs']['module_path'] == 'my.module'
|
|
140
|
+
assert result[0]['kwargs']['class_name'] == 'MyClass'
|
|
141
|
+
assert result[0]['kwargs']['method_name'] == 'my_method'
|
|
142
|
+
|
|
143
|
+
@patch('taskqueue.libs.helper_test.get_queued_tasks')
|
|
144
|
+
def test_is_task_in_queue_given_task_exists_expect_true(self, mock_get_queued_tasks):
|
|
145
|
+
mock_get_queued_tasks.return_value = [
|
|
146
|
+
{'full_func_name': 'my.module.my_function'},
|
|
147
|
+
{'full_func_name': 'other.module.other_function'}
|
|
148
|
+
]
|
|
149
|
+
|
|
150
|
+
result = is_task_in_queue('my.module.my_function')
|
|
151
|
+
|
|
152
|
+
assert result is True
|
|
153
|
+
|
|
154
|
+
@patch('taskqueue.libs.helper_test.get_queued_tasks')
|
|
155
|
+
def test_is_task_in_queue_given_task_not_exists_expect_false(self, mock_get_queued_tasks):
|
|
156
|
+
mock_get_queued_tasks.return_value = [
|
|
157
|
+
{'full_func_name': 'other.module.other_function'}
|
|
158
|
+
]
|
|
159
|
+
|
|
160
|
+
result = is_task_in_queue('my.module.my_function')
|
|
161
|
+
|
|
162
|
+
assert result is False
|
|
163
|
+
|
|
164
|
+
@patch('taskqueue.libs.helper_test.get_queued_tasks')
|
|
165
|
+
def test_is_task_in_queue_given_custom_channel_expect_channel_used(self, mock_get_queued_tasks):
|
|
166
|
+
mock_get_queued_tasks.return_value = [
|
|
167
|
+
{'full_func_name': 'my.module.my_function'}
|
|
168
|
+
]
|
|
169
|
+
|
|
170
|
+
result = is_task_in_queue('my.module.my_function', channel='high')
|
|
171
|
+
|
|
172
|
+
assert result is True
|
|
173
|
+
mock_get_queued_tasks.assert_called_once_with('high')
|
|
174
|
+
|
|
175
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
176
|
+
def test_assert_task_in_queue_given_task_exists_expect_no_exception(self, mock_is_task_in_queue):
|
|
177
|
+
mock_is_task_in_queue.return_value = True
|
|
178
|
+
|
|
179
|
+
assert_task_in_queue('my.module.my_function')
|
|
180
|
+
|
|
181
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
182
|
+
def test_assert_task_in_queue_given_task_not_exists_expect_assertion_error(self, mock_is_task_in_queue):
|
|
183
|
+
mock_is_task_in_queue.return_value = False
|
|
184
|
+
|
|
185
|
+
with pytest.raises(AssertionError, match="Task 'my.module.my_function' not found in queue 'default'"):
|
|
186
|
+
assert_task_in_queue('my.module.my_function')
|
|
187
|
+
|
|
188
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
189
|
+
def test_assert_task_in_queue_given_custom_message_expect_custom_error(self, mock_is_task_in_queue):
|
|
190
|
+
mock_is_task_in_queue.return_value = False
|
|
191
|
+
|
|
192
|
+
with pytest.raises(AssertionError, match="Custom error message"):
|
|
193
|
+
assert_task_in_queue('my.module.my_function',
|
|
194
|
+
msg="Custom error message")
|
|
195
|
+
|
|
196
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
197
|
+
def test_assert_task_in_queue_given_custom_channel_expect_channel_used(self, mock_is_task_in_queue):
|
|
198
|
+
mock_is_task_in_queue.return_value = True
|
|
199
|
+
|
|
200
|
+
assert_task_in_queue('my.module.my_function', channel='high')
|
|
201
|
+
|
|
202
|
+
mock_is_task_in_queue.assert_called_once_with(
|
|
203
|
+
'my.module.my_function', 'high')
|
|
204
|
+
|
|
205
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
206
|
+
def test_assert_task_not_in_queue_given_task_not_exists_expect_no_exception(self, mock_is_task_in_queue):
|
|
207
|
+
mock_is_task_in_queue.return_value = False
|
|
208
|
+
|
|
209
|
+
assert_task_not_in_queue('my.module.my_function')
|
|
210
|
+
|
|
211
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
212
|
+
def test_assert_task_not_in_queue_given_task_exists_expect_assertion_error(self, mock_is_task_in_queue):
|
|
213
|
+
mock_is_task_in_queue.return_value = True
|
|
214
|
+
|
|
215
|
+
with pytest.raises(AssertionError, match="Task 'my.module.my_function' found in queue 'default'"):
|
|
216
|
+
assert_task_not_in_queue('my.module.my_function')
|
|
217
|
+
|
|
218
|
+
@patch('taskqueue.libs.helper_test.is_task_in_queue')
|
|
219
|
+
def test_assert_task_not_in_queue_given_custom_message_expect_custom_error(self, mock_is_task_in_queue):
|
|
220
|
+
mock_is_task_in_queue.return_value = True
|
|
221
|
+
|
|
222
|
+
with pytest.raises(AssertionError, match="Custom error message"):
|
|
223
|
+
assert_task_not_in_queue(
|
|
224
|
+
'my.module.my_function', msg="Custom error message")
|
|
225
|
+
|
|
226
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
227
|
+
@patch('taskqueue.libs.helper_test.loads')
|
|
228
|
+
def test_celery_worker_burst_given_matching_function_expect_execution(self, mock_loads, mock_current_app):
|
|
229
|
+
mock_task = MagicMock()
|
|
230
|
+
mock_current_app.tasks = {
|
|
231
|
+
'taskqueue.cmanager.dynamic_function_executor': mock_task
|
|
232
|
+
}
|
|
233
|
+
|
|
234
|
+
mock_message = MagicMock()
|
|
235
|
+
mock_message.headers = {
|
|
236
|
+
'task': 'taskqueue.cmanager.dynamic_function_executor'}
|
|
237
|
+
mock_message.body = b'mock_body'
|
|
238
|
+
mock_message.content_type = 'application/json'
|
|
239
|
+
mock_message.content_encoding = 'utf-8'
|
|
240
|
+
mock_message.acknowledged = False
|
|
241
|
+
|
|
242
|
+
mock_loads.return_value = [[], {
|
|
243
|
+
'module_path': 'my.module',
|
|
244
|
+
'function_name': 'my_function',
|
|
245
|
+
'args': [],
|
|
246
|
+
'kwargs': {}
|
|
247
|
+
}]
|
|
248
|
+
|
|
249
|
+
mock_queue = MagicMock()
|
|
250
|
+
mock_queue.get.side_effect = [mock_message, None]
|
|
251
|
+
mock_current_app.amqp.queues = {
|
|
252
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
253
|
+
|
|
254
|
+
mock_conn = MagicMock()
|
|
255
|
+
mock_chan = MagicMock()
|
|
256
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
257
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
258
|
+
|
|
259
|
+
with patch('taskqueue.libs.helper_test.logger') as mock_logger:
|
|
260
|
+
celery_worker_burst(['my.module.my_function'])
|
|
261
|
+
|
|
262
|
+
mock_logger.info.assert_any_call(
|
|
263
|
+
"Executing task: my.module.my_function")
|
|
264
|
+
mock_logger.info.assert_any_call(
|
|
265
|
+
"Successfully executed task: my.module.my_function")
|
|
266
|
+
|
|
267
|
+
mock_message.ack.assert_called_once()
|
|
268
|
+
mock_task.apply.assert_called_once()
|
|
269
|
+
|
|
270
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
271
|
+
@patch('taskqueue.libs.helper_test.loads')
|
|
272
|
+
def test_celery_worker_burst_given_non_matching_function_expect_skip(self, mock_loads, mock_current_app):
|
|
273
|
+
mock_task = MagicMock()
|
|
274
|
+
mock_current_app.tasks = {
|
|
275
|
+
'taskqueue.cmanager.dynamic_function_executor': mock_task
|
|
276
|
+
}
|
|
277
|
+
|
|
278
|
+
mock_message = MagicMock()
|
|
279
|
+
mock_message.headers = {
|
|
280
|
+
'task': 'taskqueue.cmanager.dynamic_function_executor'}
|
|
281
|
+
mock_message.body = b'mock_body'
|
|
282
|
+
mock_message.content_type = 'application/json'
|
|
283
|
+
mock_message.content_encoding = 'utf-8'
|
|
284
|
+
mock_message.acknowledged = False
|
|
285
|
+
|
|
286
|
+
mock_loads.return_value = [[], {
|
|
287
|
+
'module_path': 'other.module',
|
|
288
|
+
'function_name': 'other_function',
|
|
289
|
+
'args': [],
|
|
290
|
+
'kwargs': {}
|
|
291
|
+
}]
|
|
292
|
+
|
|
293
|
+
mock_queue = MagicMock()
|
|
294
|
+
mock_queue.get.side_effect = [mock_message, None]
|
|
295
|
+
mock_current_app.amqp.queues = {
|
|
296
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
297
|
+
|
|
298
|
+
mock_conn = MagicMock()
|
|
299
|
+
mock_chan = MagicMock()
|
|
300
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
301
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
302
|
+
|
|
303
|
+
celery_worker_burst(['my.module.my_function'])
|
|
304
|
+
|
|
305
|
+
mock_message.ack.assert_called_once()
|
|
306
|
+
mock_task.apply.assert_not_called()
|
|
307
|
+
|
|
308
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
309
|
+
def test_celery_worker_burst_given_invalid_task_expect_warning_and_skip(self, mock_current_app):
|
|
310
|
+
mock_current_app.tasks = {}
|
|
311
|
+
|
|
312
|
+
mock_message = MagicMock()
|
|
313
|
+
mock_message.headers = {'task': 'invalid.task.name'}
|
|
314
|
+
mock_message.body = b'mock_body'
|
|
315
|
+
mock_message.content_type = 'application/json'
|
|
316
|
+
mock_message.content_encoding = 'utf-8'
|
|
317
|
+
mock_message.acknowledged = False
|
|
318
|
+
|
|
319
|
+
mock_queue = MagicMock()
|
|
320
|
+
mock_queue.get.side_effect = [mock_message, None]
|
|
321
|
+
mock_current_app.amqp.queues = {
|
|
322
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
323
|
+
|
|
324
|
+
mock_conn = MagicMock()
|
|
325
|
+
mock_chan = MagicMock()
|
|
326
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
327
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
328
|
+
|
|
329
|
+
with patch('taskqueue.libs.helper_test.logger') as mock_logger:
|
|
330
|
+
celery_worker_burst(['some.function'])
|
|
331
|
+
|
|
332
|
+
mock_logger.warning.assert_called_once_with(
|
|
333
|
+
"Invalid task 'invalid.task.name'. Skipping.")
|
|
334
|
+
mock_message.ack.assert_called_once()
|
|
335
|
+
|
|
336
|
+
@patch('taskqueue.libs.helper_test.current_app')
|
|
337
|
+
def test_celery_worker_burst_given_no_messages_expect_no_processing(self, mock_current_app):
|
|
338
|
+
mock_queue = MagicMock()
|
|
339
|
+
mock_queue.get.return_value = None
|
|
340
|
+
mock_current_app.amqp.queues = {
|
|
341
|
+
'default': MagicMock(return_value=mock_queue)}
|
|
342
|
+
|
|
343
|
+
mock_conn = MagicMock()
|
|
344
|
+
mock_chan = MagicMock()
|
|
345
|
+
mock_current_app.connection_for_read.return_value.__enter__.return_value = mock_conn
|
|
346
|
+
mock_conn.channel.return_value.__enter__.return_value = mock_chan
|
|
347
|
+
|
|
348
|
+
celery_worker_burst(['some.function'])
|
|
349
|
+
|
|
350
|
+
mock_queue.get.assert_called_once_with(no_ack=False)
|
|
@@ -1,93 +0,0 @@
|
|
|
1
|
-
import logging
|
|
2
|
-
from typing import List
|
|
3
|
-
|
|
4
|
-
from celery import current_app
|
|
5
|
-
from kombu.serialization import loads
|
|
6
|
-
|
|
7
|
-
logger = logging.getLogger(__name__)
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
def clear_all_celery_queues():
|
|
11
|
-
app = current_app
|
|
12
|
-
all_queue_names = list(app.amqp.queues.keys())
|
|
13
|
-
with app.connection_for_read() as conn:
|
|
14
|
-
with conn.channel() as chan:
|
|
15
|
-
for queue_name in all_queue_names:
|
|
16
|
-
queue = app.amqp.queues[queue_name](chan)
|
|
17
|
-
queue.purge()
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
def celery_worker_burst(include_func_names: List[str], channel: str = "default"):
|
|
21
|
-
# This doesn't use celery as celery doesn't support filtering out functions
|
|
22
|
-
# this use kombu to get the message from the queue and then execute the task manually
|
|
23
|
-
app = current_app
|
|
24
|
-
included_set = set(include_func_names)
|
|
25
|
-
processed_count = 0
|
|
26
|
-
executed_count = 0
|
|
27
|
-
|
|
28
|
-
try:
|
|
29
|
-
with app.connection_for_read() as conn:
|
|
30
|
-
with conn.channel() as chan:
|
|
31
|
-
queue = app.amqp.queues[channel](chan)
|
|
32
|
-
|
|
33
|
-
while True:
|
|
34
|
-
message = queue.get(no_ack=False)
|
|
35
|
-
if not message:
|
|
36
|
-
break
|
|
37
|
-
|
|
38
|
-
processed_count += 1
|
|
39
|
-
task_name = message.headers.get("task")
|
|
40
|
-
|
|
41
|
-
if not task_name or task_name not in app.tasks:
|
|
42
|
-
# task is not registered in celery
|
|
43
|
-
logger.warning(
|
|
44
|
-
f"Invalid task '{task_name}'. Skipping.")
|
|
45
|
-
message.ack()
|
|
46
|
-
continue
|
|
47
|
-
|
|
48
|
-
try:
|
|
49
|
-
task_obj = app.tasks[task_name]
|
|
50
|
-
accept = {"application/json",
|
|
51
|
-
"application/x-python-serialize"}
|
|
52
|
-
decoded_body = loads(
|
|
53
|
-
message.body, message.content_type, message.content_encoding, accept=accept
|
|
54
|
-
)
|
|
55
|
-
|
|
56
|
-
task_args = decoded_body[0] if decoded_body else []
|
|
57
|
-
task_kwargs = decoded_body[1] if len(
|
|
58
|
-
decoded_body) > 1 else {}
|
|
59
|
-
|
|
60
|
-
full_func_name = ""
|
|
61
|
-
if task_name.endswith("dynamic_function_executor"):
|
|
62
|
-
module_path = task_kwargs.get('module_path', '')
|
|
63
|
-
function_name = task_kwargs.get('function_name', '')
|
|
64
|
-
full_func_name = f"{module_path}.{function_name}"
|
|
65
|
-
elif task_name.endswith("dynamic_class_method_executor"):
|
|
66
|
-
module_path = task_kwargs.get('module_path', '')
|
|
67
|
-
class_name = task_kwargs.get('class_name', '')
|
|
68
|
-
method_name = task_kwargs.get('method_name', '')
|
|
69
|
-
full_func_name = f"{module_path}.{class_name}.{method_name}"
|
|
70
|
-
|
|
71
|
-
should_execute = full_func_name in included_set if full_func_name else False
|
|
72
|
-
|
|
73
|
-
if should_execute:
|
|
74
|
-
logger.info(f"Executing task: {full_func_name}")
|
|
75
|
-
message.ack()
|
|
76
|
-
task_obj.apply(args=task_args, kwargs=task_kwargs)
|
|
77
|
-
executed_count += 1
|
|
78
|
-
logger.info(
|
|
79
|
-
f"Successfully executed task: {full_func_name}")
|
|
80
|
-
else:
|
|
81
|
-
logger.info(
|
|
82
|
-
f"Skipping: {full_func_name or task_name}")
|
|
83
|
-
message.ack()
|
|
84
|
-
|
|
85
|
-
except Exception as e:
|
|
86
|
-
logger.error(
|
|
87
|
-
f"Failed to process task {task_name}: {type(e).__name__}: {e}")
|
|
88
|
-
if message and not message.acknowledged:
|
|
89
|
-
message.ack()
|
|
90
|
-
|
|
91
|
-
except Exception as e:
|
|
92
|
-
logger.error(
|
|
93
|
-
f"Failed to connect to queue {channel}: {type(e).__name__}: {e}")
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|