atagen

matrix-rust-sdk

public
Star0
Watch0
ssh://brewery@pls.lobotomise.me:33/atagen/matrix-rust-sdk.git
contrib/mitmproxy/failures.pyPython61 lines1.5 KB
1
"""
2
A mitmproxy script that introduces certain request failures in a deterministic
3
way.
4
5
Used mainly for Matrix style requests.
6
7
To run execute it with mitmproxy:
8
9
    >>> mitmproxy -s failures.py`
10
11
"""
12
import time
13
import json
14
import random
15
16
from mitmproxy import http
17
from mitmproxy.script import concurrent
18
19
REQUEST_COUNT = 0
20
21
22
def timeout(flow):
23
    timeout = 60 if "sync" in flow.request.pretty_url else 30
24
    time.sleep(timeout)
25
    return None
26
27
28
# A map holding our failure modes.
29
# The keys are just descriptive names for the failure mode while the values
30
# hold a tuple containing a function that may or may not create a failure and
31
# the probability weight at which rate this failure should be triggered.
32
#
33
# The method should return an http.Response if it should modify the
34
# response or None if the response should be passed as is.
35
FAILURES = {
36
    "Success": (lambda x: None, 50),
37
    "Gateway error":
38
    (lambda _: http.Response.make(500, b"Gateway error"), 20),
39
    "Limit exceeded": (lambda _: http.Response.make(
40
        429,
41
        json.dumps({
42
            "errcode": "M_LIMIT_EXCEEDED",
43
            "error": "Too many requests",
44
            "retry_after_ms": 2000
45
        })), 20),
46
    "Timeout error": (timeout, 10)
47
}
48
49
50
@concurrent
51
def request(flow):
52
    global FAILURES
53
54
    weights = [weight for (_, weight) in FAILURES.values()]
55
    failure = random.choices(list(FAILURES), weights=weights)[0]
56
    failure_func, _ = FAILURES[failure]
57
58
    response = failure_func(flow)
59
60
    if response:
61
        flow.response = response