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
|
# -*- coding: utf-8 -*-
"""
flask.testing
~~~~~~~~~~~~~
Implements test support helpers. This module is lazily imported
and usually not used in production environments.
:copyright: (c) 2010 by Armin Ronacher.
:license: BSD, see LICENSE for more details.
"""
from werkzeug import Client, EnvironBuilder
from flask import _request_ctx_stack
class FlaskClient(Client):
"""Works like a regular Werkzeug test client but has some
knowledge about how Flask works to defer the cleanup of the
request context stack to the end of a with body when used
in a with statement.
"""
preserve_context = context_preserved = False
def open(self, *args, **kwargs):
if self.context_preserved:
_request_ctx_stack.pop()
self.context_preserved = False
kwargs.setdefault('environ_overrides', {}) \
['flask._preserve_context'] = self.preserve_context
as_tuple = kwargs.pop('as_tuple', False)
buffered = kwargs.pop('buffered', False)
follow_redirects = kwargs.pop('follow_redirects', False)
builder = EnvironBuilder(*args, **kwargs)
if self.application.config.get('SERVER_NAME'):
server_name = self.application.config.get('SERVER_NAME')
if ':' not in server_name:
http_host, http_port = server_name, None
else:
http_host, http_port = server_name.split(':', 1)
if builder.base_url == 'http://localhost/':
# Default Generated Base URL
if http_port != None:
builder.host = http_host + ':' + http_port
else:
builder.host = http_host
old = _request_ctx_stack.top
try:
return Client.open(self, builder,
as_tuple=as_tuple,
buffered=buffered,
follow_redirects=follow_redirects)
finally:
self.context_preserved = _request_ctx_stack.top is not old
def __enter__(self):
self.preserve_context = True
return self
def __exit__(self, exc_type, exc_value, tb):
self.preserve_context = False
if self.context_preserved:
_request_ctx_stack.pop()
|
