-
-
Notifications
You must be signed in to change notification settings - Fork 41
/
pytest_socket.py
164 lines (123 loc) · 5.02 KB
/
pytest_socket.py
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
# -*- coding: utf-8 -*-
import socket
import pytest
_true_socket = socket.socket
_true_connect = socket.socket.connect
class SocketBlockedError(RuntimeError):
def __init__(self, *args, **kwargs):
super(SocketBlockedError, self).__init__("A test tried to use socket.socket.")
class SocketConnectBlockedError(RuntimeError):
def __init__(self, allowed, host, *args, **kwargs):
if allowed:
allowed = ','.join(allowed)
super(SocketConnectBlockedError, self).__init__(
'A test tried to use socket.socket.connect() with host "{0}" (allowed: "{1}").'.format(host, allowed)
)
def pytest_addoption(parser):
group = parser.getgroup('socket')
group.addoption(
'--disable-socket',
action='store_true',
dest='disable_socket',
help='Disable socket.socket by default to block network calls.'
)
group.addoption(
'--allow-hosts',
dest='allow_hosts',
metavar='ALLOWED_HOSTS_CSV',
help='Only allow specified hosts through socket.socket.connect((host, port)).'
)
group.addoption(
'--allow-unix-socket',
action='store_true',
dest='allow_unix_socket',
help='Allow calls if they are to Unix domain sockets'
)
@pytest.fixture(autouse=True)
def _socket_marker(request):
"""
Create an automatically-used fixture that every test function will load.
The last option to set the fixture wins priority.
The expected behavior is that higher granularity options should override
lower granularity options.
"""
if request.config.__socket_disabled:
request.getfixturevalue('socket_disabled')
if request.node.get_closest_marker('disable_socket'):
request.getfixturevalue('socket_disabled')
if request.node.get_closest_marker('enable_socket'):
request.getfixturevalue('socket_enabled')
@pytest.fixture
def socket_disabled(pytestconfig):
""" disable socket.socket for duration of this test function """
disable_socket(allow_unix_socket=pytestconfig.__socket_allow_unix_socket)
yield
@pytest.fixture
def socket_enabled(pytestconfig):
""" enable socket.socket for duration of this test function """
enable_socket()
yield
def disable_socket(allow_unix_socket=False):
""" disable socket.socket to disable the Internet. useful in testing.
"""
class GuardedSocket(socket.socket):
""" socket guard to disable socket creation (from pytest-socket) """
def __new__(cls, family=-1, type=-1, proto=-1, fileno=None):
try:
is_unix_socket = family == socket.AF_UNIX
except AttributeError:
# AF_UNIX not supported on Windows https://bugs.python.org/issue33408
is_unix_socket = False
if is_unix_socket and allow_unix_socket:
return super().__new__(cls, family, type, proto, fileno)
raise SocketBlockedError()
socket.socket = GuardedSocket
def enable_socket():
""" re-enable socket.socket to enable the Internet. useful in testing.
"""
socket.socket = _true_socket
def pytest_configure(config):
config.addinivalue_line("markers", "disable_socket(): Disable socket connections for a specific test")
config.addinivalue_line("markers", "enable_socket(): Enable socket connections for a specific test")
config.addinivalue_line("markers", "allow_hosts([hosts]): Restrict socket connection to defined list of hosts")
# Store the global configs in the `pytest.Config` object.
config.__socket_disabled = config.getoption('--disable-socket')
config.__socket_allow_unix_socket = config.getoption('--allow-unix-socket')
config.__socket_allow_hosts = config.getoption('--allow-hosts')
def pytest_runtest_setup(item):
mark_restrictions = item.get_closest_marker('allow_hosts')
cli_restrictions = item.config.__socket_allow_hosts
hosts = None
if mark_restrictions:
hosts = mark_restrictions.args[0]
elif cli_restrictions:
hosts = cli_restrictions
socket_allow_hosts(hosts)
def pytest_runtest_teardown():
_remove_restrictions()
def host_from_address(address):
host = address[0]
if isinstance(host, str):
return host
def host_from_connect_args(args):
address = args[0]
if isinstance(address, tuple):
return host_from_address(address)
def socket_allow_hosts(allowed=None):
""" disable socket.socket.connect() to disable the Internet. useful in testing.
"""
if isinstance(allowed, str):
allowed = allowed.split(',')
if not isinstance(allowed, list):
return
def guarded_connect(inst, *args):
host = host_from_connect_args(args)
if host and host in allowed:
return _true_connect(inst, *args)
raise SocketConnectBlockedError(allowed, host)
socket.socket.connect = guarded_connect
def _remove_restrictions():
""" restore socket.socket.* to allow access to the Internet. useful in testing.
"""
socket.socket = _true_socket
socket.socket.connect = _true_connect