Added IPV4 and IPV6 only fixtures

azmeuk · csernazs · commit 155918bb16a1 · 2022-09-24T21:53:17.000+02:00
diff --git a/pytest_httpserver/httpserver.py b/pytest_httpserver/httpserver.py
@@ -1,4 +1,5 @@
 import abc
+import ipaddress
 import json
 import queue
 import re
@@ -663,7 +664,9 @@ def url_for(self, suffix: str):
         else:
             protocol = "https"
 
-        return "{}://{}:{}{}".format(protocol, self.host, self.port, suffix)
+        host = self.format_host(self.host)
+
+        return "{}://{}:{}{}".format(protocol, host, self.port, suffix)
 
     def create_matcher(self, *args, **kwargs) -> RequestMatcher:
         """
@@ -837,6 +840,24 @@ def __exit__(self, *args, **kwargs):
         if self.is_running():
             self.stop()
 
+    @staticmethod
+    def format_host(host):
+        """
+        Formats a hostname so it can be used in a URL.
+        Notably, this adds brackets around IPV6 addresses when
+        they are missing.
+        """
+        try:
+            ipaddress.IPv6Address(host)
+            is_ipv6 = True
+        except ValueError:
+            is_ipv6 = False
+
+        if is_ipv6 and not host.startswith("[") and not host.endswith("]"):
+            return f"[{host}]"
+
+        return host
+
 
 class HTTPServer(HTTPServerBase):  # pylint: disable=too-many-instance-attributes
     """
diff --git a/pytest_httpserver/pytest_plugin.py b/pytest_httpserver/pytest_plugin.py
@@ -66,3 +66,37 @@ def httpserver(make_httpserver):
     server = make_httpserver
     yield server
     server.clear()
+
+
+@pytest.fixture(scope="session")
+def make_httpserver_ipv4(httpserver_ssl_context):
+    server = HTTPServer(host="127.0.0.1", port=0, ssl_context=httpserver_ssl_context)
+    server.start()
+    yield server
+    server.clear()
+    if server.is_running():
+        server.stop()
+
+
+@pytest.fixture
+def httpserver_ipv4(make_httpserver_ipv4):
+    server = make_httpserver_ipv4
+    yield server
+    server.clear()
+
+
+@pytest.fixture(scope="session")
+def make_httpserver_ipv6(httpserver_ssl_context):
+    server = HTTPServer(host="::1", port=0, ssl_context=httpserver_ssl_context)
+    server.start()
+    yield server
+    server.clear()
+    if server.is_running():
+        server.stop()
+
+
+@pytest.fixture
+def httpserver_ipv6(make_httpserver_ipv6):
+    server = make_httpserver_ipv6
+    yield server
+    server.clear()
diff --git a/tests/test_ip_protocols.py b/tests/test_ip_protocols.py
@@ -0,0 +1,18 @@
+import requests
+
+
+def test_ipv4(httpserver_ipv4):
+    httpserver_ipv4.expect_request("/").respond_with_data("OK")
+    assert httpserver_ipv4.host == "127.0.0.1"
+
+    response = requests.get(httpserver_ipv4.url_for("/"))
+    assert response.text == "OK"
+
+
+def test_ipv6(httpserver_ipv6):
+    httpserver_ipv6.expect_request("/").respond_with_data("OK")
+    assert httpserver_ipv6.host == "::1"
+    assert httpserver_ipv6.url_for("/") == f"http://[::1]:{httpserver_ipv6.port}/"
+
+    response = requests.get(httpserver_ipv6.url_for("/"))
+    assert response.text == "OK"
diff --git a/tests/test_release.py b/tests/test_release.py
@@ -199,6 +199,7 @@ def test_sdist_contents(build: Build, version: str):
             "test_blocking_httpserver_howto.py",
             "test_handler_errors.py",
             "test_headers.py",
+            "test_ip_protocols.py",
             "test_json_matcher.py",
             "test_mixed.py",
             "test_oneshot.py",
