Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Disable the internet in Python. With py.test hooks. (Disable socket.socket.) GREAT for unit testing.
from __future__ import print_function
import socket
import sys
_module = sys.modules[__name__]
def disable_socket():
""" disable socket.socket to disable the Internet. useful in testing.
.. doctest::
>>> # we'll try httplib a little later, but import it immediately, before tampering.
>>> # hopefully this proves that the 'patch' doesn't have to happen before other imports.
>>> import httplib
>>> enable_socket() # should be able to call "enable" at any time, even when enabled...
[!] socket.socket is UN-blocked, and the network can be accessed.
>>> disable_socket() # OK let's disable it.
[!] socket.socket is now blocked. The network should be inaccessible.
>>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Traceback (most recent call last):
...
RuntimeError: A test tried to use socket.socket without explicitly un-blocking it.
>>> httplib.HTTPConnection("scanme.nmap.org:80").request("GET", "/")
Traceback (most recent call last):
...
RuntimeError: A test tried to use socket.socket without explicitly un-blocking it.
>>> enable_socket()
[!] socket.socket is UN-blocked, and the network can be accessed.
>>> enable_socket() # twice in a row should work.
[!] socket.socket is UN-blocked, and the network can be accessed.
>>> disable_socket()
[!] socket.socket is now blocked. The network should be inaccessible.
>>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Traceback (most recent call last):
...
RuntimeError: A test tried to use socket.socket without explicitly un-blocking it.
>>> enable_socket()
[!] socket.socket is UN-blocked, and the network can be accessed.
"""
setattr(_module, u'_socket_disabled', True)
def guarded(*args, **kwargs):
if getattr(_module, u'_socket_disabled', False):
raise RuntimeError(
u"A test tried to use socket.socket without explicitly un-blocking it.")
else:
# SocketType is a valid, public alias of socket.socket,
# we use it here to avoid namespace collisions
return socket.SocketType(*args, **kwargs)
socket.socket = guarded
print(u'[!] socket.socket is now blocked. The network should be inaccessible.')
def enable_socket():
""" re-enable socket.socket to enable the Internet. useful in testing.
"""
setattr(_module, u'_socket_disabled', False)
print(u'[!] socket.socket is UN-blocked, and the network can be accessed.')
# Put this in the conftest.py at the top of your unit tests folder,
# so it's available to all unit tests
import pytest
import _socket_toggle
def pytest_runtest_setup():
""" disable the interet. test-cases can explicitly re-enable """
_socket_toggle.disable_socket()
@pytest.fixture(scope='function')
def enable_socket(request):
""" re-enable socket.socket for duration of this test function """
_socket_toggle.enable_socket()
request.addfinalizer(_socket_toggle.disable_socket)
# Example usage of the py.test fixture in tests
import socket
import pytest
try:
from urllib2 import urlopen
except ImportError:
import urllib3
urlopen = urllib.request.urlopen
def test_socket_disabled_by_default():
# default behavior: socket.socket is unusable
with pytest.raises(RuntimeError):
urlopen(u'https://www.python.org/')
def test_explicitly_enable_socket(enable_socket):
# socket is enabled by pytest fixture from conftest. disabled in finalizer
assert socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@hangtwenty

This comment has been minimized.

@hangtwenty

This comment has been minimized.

Copy link
Owner Author

@hangtwenty hangtwenty commented Jun 1, 2017

Use the real pytest plugin instead!

This has been made into a real pytest plugin.

https://github.com/miketheman/pytest-socket

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.