Python: 为了测试目的，阻止网络连接？

Question

我正在尝试测试一个提供几个网络服务接口的软件包。它有一个测试套件，可以在大多数情况下测试功能，而无需连接到互联网。然而，还有一些残留的测试可能会尝试连接到互联网/下载数据，我希望阻止它们这样做有两个原因：首先，确保我的测试套件在没有网络连接的情况下工作；其次，避免向网络服务发送过多的查询。

一个明显的解决方案是拔掉我的机器/关闭无线连接，但当我在远程机器上运行测试时，这显然行不通。

所以，我的问题是：我能够阻止单个Python进程访问网络/端口吗？（"沙盒化"它，只是阻止网络连接）

（据我所知，pysandbox不能做到这一点）

我正在使用py.test，所以我需要一个能与py.test配合使用的解决方案，以防对任何提出的答案产生影响。

- keflavich

---

对于那些可能正在寻找可定制的阻止和/或记录连接的人，最好使用vcrpy。它有一个pytest插件。 - shad0w_wa1k3r

Answer 1

使用“Monkey patching”socket 应该可以解决此问题：

import socket
def guard(*args, **kwargs):
    raise Exception("I told you not to use the Internet!")
socket.socket = guard

请确保此代码在任何其他引用之前运行。

- Thomas Orozco

---

这太棒了！你有什么想法如何让py.test在其他任何操作之前运行它吗？ - keflavich

回复我的最后一条评论：在 conftests.py 中运行此代码。 - keflavich

现在有一个 Py.test 插件，因此尽可能使用它。如果无法使用，则可以在 socket 方法上使用 patch，并将 side_effect=Exception 作为参数传递。 - Pieter

这是一个不错的解决方案，但值得注意的是它只影响使用Python socket API的代码。直接调用系统的代码（例如封装C库的模块）不受此影响。 - el.pescado - нет войне

这些解决方案都不能适用于许多情况，甚至似乎无法阻止请求库的使用。 - Steve Yeago

Answer 2

更新：现在有一个 pytest 插件可以完成与本答案相同的任务！你可以阅读这个答案了解其工作原理，但是我强烈建议使用插件而不是复制粘贴我的答案 :-) 点击这里查看：https://github.com/miketheman/pytest-socket

---

我发现 Thomas Orozco 的答案非常有帮助。继 keflavich 之后，这是如何将其集成到我的单元测试套件中的方法。这适用于我具有数千个非常不同的单元测试用例（<100 需要套接字）... 还在和不在 doctest 中。

我在这里发布了它。下面包括方便起见。使用 Python 2.7.5、pytest==2.7.0 进行测试。（要自己测试，请在克隆了所有 3 个文件的目录中运行py.test --doctest-modules。）

_socket_toggle.py

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::
        >>> enable_socket()
        [!] socket.socket is enabled.
        >>> disable_socket()
        [!] socket.socket is disabled. Welcome to the desert of the real.
        >>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        Traceback (most recent call last):
        ...
        RuntimeError: I told you not to use the Internet!
        >>> enable_socket()
        [!] socket.socket is enabled.
        >>> enable_socket()
        [!] socket.socket is enabled.
        >>> disable_socket()
        [!] socket.socket is disabled. Welcome to the desert of the real.
        >>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        Traceback (most recent call last):
        ...
        RuntimeError: I told you not to use the Internet!
        >>> enable_socket()
        [!] socket.socket is enabled.
    """
    setattr(_module, '_socket_disabled', True)

    def guarded(*args, **kwargs):
        if getattr(_module, '_socket_disabled', False):
            raise RuntimeError("I told you not to use the Internet!")
        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 disabled. Welcome to the desert of the real.')


def enable_socket():
    """ re-enable socket.socket to enable the Internet. useful in testing.
    """
    setattr(_module, '_socket_disabled', False)
    print(u'[!] socket.socket is enabled.')

conftest.py

# 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)

test_example.py

# 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)

- floer32

---

为什么不抛出 ConnectionError 异常？ - Femto Trader

@FemtoTrader 我猜应该是因为那个太正确了。不想混淆我们抛出的错误和内置错误抛出的错误，对于一个合法的 ConnectionError。事实上我实际上使用了一个 Runtime Error 的子类，但是我想让这个例子更简单。 - floer32

有没有更新的版本适用于Python3？我也很乐意帮忙把这个变成pytest插件。 - Mike Fiedler

@MikeFiedler 我已经更新了答案，链接到了你的插件。做得好！ - floer32

太棒了。根据我们设置的测试方式，这是一个很好的额外检查。 - Oliver Shaw

Answer 3

在Thomas Orozco和driftcatcher非常有帮助的答案基础上，这是一个与Python的unittest和Django兼容的变体。

您只需要从增强的NoSocketTestCase类继承您的测试用例类，并且任何对网络的访问都将被检测到并引发SocketAccessError异常。

而且，这种方法也适用于Django。您只需要将NoSocketTestCase类更改为从django.test.TestCase而不是unittest.TestCase继承即可。

虽然它并没有严格回答OP的问题，但我认为这可能对任何想要在单元测试中阻止网络访问的人有所帮助。

no_sockets.py

import socket
from unittest import TestCase


class SocketAccessError(Exception):
    pass


class NoSocketsTestCase(TestCase):
    """Enhancement of TestCase class that prevents any use of sockets

    Will throw the exception SocketAccessError when any code tries to
    access network sockets
    """

    @classmethod
    def setUpClass(cls):
        cls.socket_original = socket.socket
        socket.socket = cls.guard
        return super().setUpClass()

    @classmethod
    def tearDownClass(cls):
        socket.socket = cls.socket_original
        return super().tearDownClass()

    @staticmethod
    def guard(*args, **kwargs):
        raise SocketAccessError('Attempted to access network')

test_no_sockets.py

import urllib.request
from .no_sockets import NoSocketsTestCase, SocketAccessError


class TestNoSocketsTestCase(NoSocketsTestCase):

    def test_raises_exception_on_attempted_network_access(self):

        with self.assertRaises(SocketAccessError):            
            urllib.request.urlopen('https://www.google.com')

- Erik Kalkoken

Answer 4

一个简单的方法来限制 requests 库:

from unittest import mock

requests_gag = mock.patch(
    'requests.Session.request',
    mock.Mock(side_effect=RuntimeError(
        'Please use the `responses` library to mock HTTP in your tests.'
    ))
)

with requests_gag:
    ...  # no Internet here

- kolypto

Answer 5

httpretty是一个解决这个问题的小型库。

如果您正在使用Django测试运行程序，请编写自定义测试运行程序，在其中禁用所有第三方API调用。

# common/test_runner.py

import httpretty
from django.test.runner import DiscoverRunner


class CustomTestRunner(DiscoverRunner):
    def run_tests(self, *args, **kwargs):
        with httpretty.enabled(allow_net_connect=False):
            return super().run_tests(*args, **kwargs)

将这个新的测试运行器添加到您的设置中。

TEST_RUNNER = "common.test_runner.CustomTestRunner"

从现在开始，所有的外部API调用都必须被模拟，否则将会引发httpretty.errors.UnmockedError。

如果您正在使用pytest，这个fixture应该可以工作。

@pytest.fixture
def disable_external_api_calls():
    httpretty.enable()
    yield
    httpretty.disable()

- Tom Wojcik

Answer 6

我有一个pytest解决方案。 pytest-network 库对我很有帮助。

# conftest.py
import pytest
import socket

_original_connect = socket.socket.connect

def patched_connect(*args, **kwargs):
    ...
    # It depends on your testing purpose
    # You may want a exception, add here
    # If you test unconnectable situations
    # it can stay like this 
    

@pytest.fixture
def enable_network():
    socket.socket.connect = _original_connect
    yield
    socket.socket.connect = patched_connect

@pytest.fixture
def disable_network():
    socket.socket.connect = patched_connect
    yield
    socket.socket.connect = _original_connect

# test_internet.py
def test_your_unconnectable_situation(disable_network):
    response = request.get('http://stackoverflow.com/')
    response.status_code == 400

- Baris Senyerli

Answer 7

其他答案会破坏只创建本地线程的测试，比如pytorch的数据加载器。如果你想允许本地主机的网络连接，你可以尝试这段代码。

    import socket
    orig_connect = socket.socket.connect
    def guard(self, address: tuple, *args: Any, **kwargs: Any) -> None:
        is_local_connection = False
        # Family types documeted at https://docs.python.org/3/library/socket.html#socket-families
        if self.family == socket.AF_UNIX:
            # AF_UNIX is for local connections only
            is_local_connection = True
        elif self.family == socket.AF_INET or self.family == socket.AF_INET6:
            # Check that the host is local
            host = address[0]
            is_local_connection = host  == 'localhost' or host == '127.0.0.1'
        else:
            # Other types are too rare to bother with (e.g. cluster computing)
            pass
        if is_local_connection:
            orig_connect(self, address, *args, **kwargs)
        else:
            assert False, "No internet allowed in unittests"
    socket.socket.connect = guard

- Alex Li