diff --git a/pytest_socket.py b/pytest_socket.py index b62e511..639bdd1 100644 --- a/pytest_socket.py +++ b/pytest_socket.py @@ -4,6 +4,8 @@ _true_socket = socket.socket _true_connect = socket.socket.connect +_true_getaddrinfo = socket.getaddrinfo +_true_gethostbyname = socket.gethostbyname class SocketBlockedError(RuntimeError): @@ -78,12 +80,16 @@ def __new__(cls, family=-1, type=-1, proto=-1, fileno=None): raise SocketBlockedError() socket.socket = GuardedSocket + socket.getaddrinfo = GuardedSocket + socket.gethostbyname = GuardedSocket def enable_socket(): """ re-enable socket.socket to enable the Internet. useful in testing. """ socket.socket = _true_socket + socket.getaddrinfo = _true_getaddrinfo + socket.gethostbyname = _true_gethostbyname def pytest_configure(config): diff --git a/tests/test_socket.py b/tests/test_socket.py index b2ec58c..827af76 100644 --- a/tests/test_socket.py +++ b/tests/test_socket.py @@ -8,6 +8,7 @@ def test_socket(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """ @@ -62,6 +63,7 @@ def disable_socket_for_all(): def test_socket(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose") assert_socket_blocked(result) @@ -117,6 +119,7 @@ def test_disable_socket_marker(testdir): @pytest.mark.disable_socket def test_socket(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose") assert_socket_blocked(result) @@ -130,6 +133,7 @@ def test_enable_socket_marker(testdir): @pytest.mark.enable_socket def test_socket(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose", "--disable-socket") result.assert_outcomes(1, 0, 0) @@ -198,17 +202,20 @@ def test_double_enabled(): pytest_socket.enable_socket() pytest_socket.enable_socket() socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") def test_double_disabled(): pytest_socket.disable_socket() pytest_socket.disable_socket() with pytest.raises(pytest_socket.SocketBlockedError): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") def test_disable_enable(): pytest_socket.disable_socket() pytest_socket.enable_socket() socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose") result.assert_outcomes(3, 0, 0) @@ -221,6 +228,7 @@ def test_socket_enabled_fixture(testdir, socket_enabled): import socket def test_socket_enabled(socket_enabled): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose") result.assert_outcomes(1, 0, 0) @@ -234,10 +242,15 @@ def test_mix_and_match(testdir, socket_enabled): def test_socket1(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") + def test_socket_enabled(socket_enabled): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") + def test_socket2(): socket.socket(socket.AF_INET, socket.SOCK_STREAM) + socket.gethostbyname("localhost") """) result = testdir.runpytest("--verbose", "--disable-socket") result.assert_outcomes(1, 0, 2)