@@ -44,40 +44,27 @@ def pytest_addoption(parser):
4444 )
4545
4646
47- @pytest .fixture (autouse = True )
48- def _socket_marker (request ):
49- """
50- Create an automatically-used fixture that every test function will load.
51-
52- The last option to set the fixture wins priority.
53- The expected behavior is that higher granularity options should override
54- lower granularity options.
55- """
56- if request .config .getoption ('--disable-socket' ):
57- request .getfixturevalue ('socket_disabled' )
58-
59- if request .node .get_closest_marker ('disable_socket' ):
60- request .getfixturevalue ('socket_disabled' )
61- if request .node .get_closest_marker ('enable_socket' ):
62- request .getfixturevalue ('socket_enabled' )
63-
64-
6547@pytest .fixture
6648def socket_disabled (pytestconfig ):
6749 """ disable socket.socket for duration of this test function """
68- allow_unix_socket = pytestconfig .getoption ('--allow-unix-socket' )
69- disable_socket (allow_unix_socket )
50+ disable_socket (allow_unix_socket = pytestconfig .__socket_allow_unix_socket )
7051 yield
71- enable_socket ()
7252
7353
7454@pytest .fixture
7555def socket_enabled (pytestconfig ):
7656 """ enable socket.socket for duration of this test function """
7757 enable_socket ()
7858 yield
79- allow_unix_socket = pytestconfig .getoption ('--allow-unix-socket' )
80- disable_socket (allow_unix_socket )
59+
60+
61+ def _is_unix_socket (family ) -> bool :
62+ try :
63+ is_unix_socket = family == socket .AF_UNIX
64+ except AttributeError :
65+ # AF_UNIX not supported on Windows https://bugs.python.org/issue33408
66+ is_unix_socket = False
67+ return is_unix_socket
8168
8269
8370def disable_socket (allow_unix_socket = False ):
@@ -86,15 +73,9 @@ def disable_socket(allow_unix_socket=False):
8673
8774 class GuardedSocket (socket .socket ):
8875 """ socket guard to disable socket creation (from pytest-socket) """
89- def __new__ (cls , * args , ** kwargs ):
90- try :
91- is_unix_socket = args [0 ] == socket .AF_UNIX
92- except AttributeError :
93- # AF_UNIX not supported on Windows https://bugs.python.org/issue33408
94- is_unix_socket = False
95-
96- if is_unix_socket and allow_unix_socket :
97- return super ().__new__ (cls , * args , ** kwargs )
76+ def __new__ (cls , family = - 1 , type = - 1 , proto = - 1 , fileno = None ):
77+ if _is_unix_socket (family ) and allow_unix_socket :
78+ return super ().__new__ (cls , family , type , proto , fileno )
9879
9980 raise SocketBlockedError ()
10081
@@ -116,20 +97,54 @@ def pytest_configure(config):
11697 config .addinivalue_line ("markers" , "enable_socket(): Enable socket connections for a specific test" )
11798 config .addinivalue_line ("markers" , "allow_hosts([hosts]): Restrict socket connection to defined list of hosts" )
11899
100+ # Store the global configs in the `pytest.Config` object.
101+ config .__socket_disabled = config .getoption ('--disable-socket' )
102+ config .__socket_allow_unix_socket = config .getoption ('--allow-unix-socket' )
103+ config .__socket_allow_hosts = config .getoption ('--allow-hosts' )
104+
105+
106+ def pytest_runtest_setup (item ) -> None :
107+ """During each test item's setup phase,
108+ choose the behavior based on the configurations supplied.
119109
120- def pytest_runtest_setup (item ):
110+ This is the bulk of the logic for the plugin.
111+ As the logic can be extensive, this method is allowed complexity.
112+ It may be refactored in the future to be more readable.
113+ """
114+
115+ # If test has the `enable_socket` marker, we accept this as most explicit.
116+ if 'socket_enabled' in item .fixturenames or item .get_closest_marker ('enable_socket' ):
117+ enable_socket ()
118+ return
119+
120+ # If the test has the `disable_socket` marker, it's explicitly disabled.
121+ if 'socket_disabled' in item .fixturenames or item .get_closest_marker ('disable_socket' ):
122+ disable_socket (item .config .__socket_allow_unix_socket )
123+ return
124+
125+ # Resolve `allow_hosts` behaviors.
126+ hosts = _resolve_allow_hosts (item )
127+
128+ # Finally, check the global config and disable socket if needed.
129+ if item .config .__socket_disabled and not hosts :
130+ disable_socket (item .config .__socket_allow_unix_socket )
131+
132+
133+ def _resolve_allow_hosts (item ):
134+ """Resolve `allow_hosts` behaviors."""
121135 mark_restrictions = item .get_closest_marker ('allow_hosts' )
122- cli_restrictions = item .config .getoption ( '--allow-hosts' )
136+ cli_restrictions = item .config .__socket_allow_hosts
123137 hosts = None
124138 if mark_restrictions :
125139 hosts = mark_restrictions .args [0 ]
126140 elif cli_restrictions :
127141 hosts = cli_restrictions
128- socket_allow_hosts (hosts )
142+ socket_allow_hosts (hosts , allow_unix_socket = item .config .__socket_allow_unix_socket )
143+ return hosts
129144
130145
131146def pytest_runtest_teardown ():
132- remove_host_restrictions ()
147+ _remove_restrictions ()
133148
134149
135150def host_from_address (address ):
@@ -145,7 +160,7 @@ def host_from_connect_args(args):
145160 return host_from_address (address )
146161
147162
148- def socket_allow_hosts (allowed = None ):
163+ def socket_allow_hosts (allowed = None , allow_unix_socket = False ):
149164 """ disable socket.socket.connect() to disable the Internet. useful in testing.
150165 """
151166 if isinstance (allowed , str ):
@@ -155,14 +170,16 @@ def socket_allow_hosts(allowed=None):
155170
156171 def guarded_connect (inst , * args ):
157172 host = host_from_connect_args (args )
158- if host and host in allowed :
173+ if host in allowed or ( _is_unix_socket ( inst . family ) and allow_unix_socket ) :
159174 return _true_connect (inst , * args )
175+
160176 raise SocketConnectBlockedError (allowed , host )
161177
162178 socket .socket .connect = guarded_connect
163179
164180
165- def remove_host_restrictions ():
166- """ restore socket.socket.connect() to allow access to the Internet. useful in testing.
181+ def _remove_restrictions ():
182+ """ restore socket.socket.* to allow access to the Internet. useful in testing.
167183 """
184+ socket .socket = _true_socket
168185 socket .socket .connect = _true_connect
0 commit comments