mrq/tests/test_parallel.py at 3fbdd2e2408b0e972c372c8b73d1acd890031908 · pricingassistant/mrq · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from builtins import range
import time
import pytest
from mrq.context import connections


@pytest.mark.parametrize(["p_flags"], [
    ["--greenlets 50"],
    ["--processes 10 --greenlets 5"]
])
def test_parallel_100sleeps(worker, p_flags):

    worker.start(flags=p_flags)

    start_time = time.time()

    # This will sleep a total of 100 seconds
    result = worker.send_tasks(
        "tests.tasks.general.Add", [{"a": i, "b": 0, "sleep": 1} for i in range(100)])

    total_time = time.time() - start_time

    # But should be done quickly!
    assert total_time < 15

    # ... and return correct results
    assert result == list(range(100))


@pytest.mark.parametrize(["p_greenlets"], [
    [1],
    [2]
])
def test_dequeue_strategy(worker, p_greenlets):

    worker.start_deps(flush=True)

    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 41, "sleep": 2}, queue="q1", block=False, start=False)
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 42, "sleep": 2}, queue="q2", block=False, start=False)
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 41, "sleep": 2}, queue="q1", block=False, start=False)
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 42, "sleep": 2}, queue="q2", block=False, start=False)
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 43, "sleep": 2}, queue="q3", block=False, start=False)
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 43, "sleep": 2}, queue="q3", block=False, start=False)

    time.sleep(0.1)

    worker.start(flags="--dequeue_strategy parallel --greenlets %s" % p_greenlets, queues="q1 q2", deps=False, start=False)

    if p_greenlets == 1:
        time.sleep(1 + 2)
    else:
        time.sleep(1)

    # Should be dequeued in parallel
    assert connections.mongodb_jobs.tests_inserts.count({"params.a": 41}) == 1
    assert connections.mongodb_jobs.tests_inserts.count({"params.a": 42}) == 1
    assert connections.mongodb_jobs.tests_inserts.count() == 2

    worker.stop(deps=False, sig=9)
    time.sleep(1)

    worker.start(flags="--dequeue_strategy burst --greenlets 2", queues="q3", deps=False)

    time.sleep(3)

    assert connections.mongodb_jobs.tests_inserts.count({"params.a": 43}) == 2

    # Worker should be stopped now so even if we queue nothing will happen.
    worker.send_task(
        "tests.tasks.general.MongoInsert", {"a": 43, "sleep": 2}, queue="q3", block=False, start=False)

    time.sleep(2)

    assert connections.mongodb_jobs.tests_inserts.count({"params.a": 43}) == 2

    worker.stop()