File: util.py

package info (click to toggle)
python-molotov 2.7-3
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 8,268 kB
  • sloc: python: 4,121; makefile: 60
file content (228 lines) | stat: -rw-r--r-- 5,627 bytes parent folder | download
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
import asyncio
import functools
import json
import os
import sys
import threading
import time
import traceback
from io import StringIO

from aiohttp import ClientSession, __version__

_DNS_CACHE = {}
_STOP = False
_STOP_WHY = []
_TIMER = None
if __version__[0] == "2":
    raise ImportError("Molotov only supports aiohttp 3.x going forward")


def event_loop():
    if sys.version_info.minor >= 10:
        try:
            return asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            return loop
    return asyncio.get_event_loop()


def get_timer():
    return _TIMER


def set_timer(value=None):
    global _TIMER
    if value is None:
        value = int(time.time())
    _TIMER = value


def stop(why=None):
    global _STOP
    if why is not None:
        _STOP_WHY.append(why)
    _STOP = True


def stop_reason():
    return _STOP_WHY


def is_stopped():
    return _STOP


class OptionError(Exception):
    pass


def _expand_args(args, options):
    for key, val in options.items():
        setattr(args, key, val)


def expand_options(config, scenario, args):
    if not isinstance(config, str):
        try:
            config = json.loads(config.read())
        except Exception as err:
            raise OptionError("Can't parse %r" % config) from err
    else:
        if not os.path.exists(config):
            raise OptionError("Can't find %r" % config)

        with open(config) as f:
            try:
                config = json.loads(f.read())
            except ValueError as err:
                raise OptionError("Can't parse %r" % config) from err

    if "molotov" not in config:
        raise OptionError("Bad config -- no molotov key")

    if "tests" not in config["molotov"]:
        raise OptionError("Bad config -- no molotov/tests key")

    if scenario not in config["molotov"]["tests"]:
        raise OptionError("Can't find %r in the config" % scenario)

    _expand_args(args, config["molotov"]["tests"][scenario])


def _run_in_fresh_loop(coro, timeout=30):
    thres = []
    thexc = []

    def run():
        loop = asyncio.new_event_loop()
        try:
            task = loop.create_task(coro())
            thres.append(loop.run_until_complete(task))
        except Exception as e:
            thexc.append(e)
        finally:
            loop.close()

    th = threading.Thread(target=run)
    th.start()
    th.join(timeout=timeout)

    # re-raise a thread exception
    if len(thexc) > 0:
        raise thexc[0]

    return thres[0]


async def _request(endpoint, verb="GET", session_options=None, json=False, **options):
    if session_options is None:
        session_options = {}

    async with ClientSession(**session_options) as session:
        meth = getattr(session, verb.lower())
        result = {}
        async with meth(endpoint, **options) as resp:
            if json:
                result["content"] = await resp.json()
            else:
                result["content"] = await resp.text()
            result["status"] = resp.status
            result["headers"] = resp.headers

        return result


def request(endpoint, verb="GET", session_options=None, **options):
    """Performs a synchronous request.

    Uses a dedicated event loop and aiohttp.ClientSession object.

    Options:

    - endpoint: the endpoint to call
    - verb: the HTTP verb to use (defaults: GET)
    - session_options: a dict containing options to initialize the session
      (defaults: None)
    - options: extra options for the request (defaults: None)

    Returns a dict object with the following keys:

    - content: the content of the response
    - status: the status
    - headers: a dict with all the response headers
    """
    req = functools.partial(_request, endpoint, verb, session_options, **options)
    return _run_in_fresh_loop(req)


def json_request(endpoint, verb="GET", session_options=None, **options):
    """Like :func:`molotov.request` but extracts json from the response."""
    req = functools.partial(_request, endpoint, verb, session_options, json=True, **options)
    return _run_in_fresh_loop(req)


_VARS = {}


def set_var(name, value):
    """Sets a global variable.

    Options:

    - name: name of the variable
    - value: object to set
    """
    _VARS[name] = value


def get_var(name, factory=None):
    """Gets a global variable given its name.

    If factory is not None and the variable is not set, factory
    is a callable that will set the variable.

    If not set, returns None.
    """
    if name not in _VARS and factory is not None:
        _VARS[name] = factory()
    return _VARS.get(name)


# taken from https://stackoverflow.com/a/37211337
def _make_sleep():
    async def sleep(delay, result=None, *, loop=None):
        coro = asyncio.sleep(delay, result=result)
        task = asyncio.ensure_future(coro, loop=loop)
        sleep.tasks.add(task)
        try:
            return await task
        except asyncio.CancelledError:
            return result
        finally:
            sleep.tasks.remove(task)

    sleep.tasks = set()
    sleep.cancel_all = lambda: sum(task.cancel() for task in sleep.tasks)
    return sleep


cancellable_sleep = _make_sleep()


def printable_error(error, tb=None):
    printable = [repr(error)]
    if tb is None:
        tb = sys.exc_info()[2]
    printed = StringIO()
    traceback.print_tb(tb, file=printed)
    printed.seek(0)
    for line in printed.readlines():
        printable.append(line.rstrip("\n"))
    return printable


def now():
    return int(time.time())