mirror of
https://github.com/esphome/esphome.git
synced 2026-08-17 02:47:52 +08:00
[core] Use Happy Eyeballs for remote file downloads (#18050)
This commit is contained in:
@@ -12,6 +12,7 @@ from esphome.components.packages import validate_source_shorthand
|
||||
import esphome.config_validation as cv
|
||||
from esphome.const import CONF_ESPHOME, CONF_PROJECT, CONF_REF, CONF_WIFI
|
||||
import esphome.final_validate as fv
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
from esphome.yaml_util import dump
|
||||
|
||||
dashboard_import_ns = cg.esphome_ns.namespace("dashboard_import")
|
||||
@@ -109,6 +110,7 @@ def import_config(
|
||||
if git_file.query and "full_config" in git_file.query:
|
||||
url = git_file.raw_url
|
||||
try:
|
||||
ensure_happy_eyeballs()
|
||||
req = requests.get(url, timeout=30)
|
||||
req.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
|
||||
@@ -3286,9 +3286,19 @@ def copy_files():
|
||||
if str(path).startswith("http"):
|
||||
import requests
|
||||
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
|
||||
try:
|
||||
req = requests.get(path, timeout=30)
|
||||
req.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise EsphomeError(
|
||||
f"Could not download extra build file {path}: {e}"
|
||||
) from e
|
||||
CORE.relative_build_path(name).parent.mkdir(parents=True, exist_ok=True)
|
||||
content = requests.get(path, timeout=30).content
|
||||
CORE.relative_build_path(name).write_bytes(content)
|
||||
CORE.relative_build_path(name).write_bytes(req.content)
|
||||
else:
|
||||
copy_file_if_changed(path, CORE.relative_build_path(name))
|
||||
|
||||
|
||||
@@ -36,6 +36,7 @@ from esphome.const import (
|
||||
CONF_WEIGHT,
|
||||
)
|
||||
from esphome.core import CORE, HexInt
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
from esphome.types import ConfigType
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
@@ -319,6 +320,7 @@ def download_gfont(value):
|
||||
if not external_files.is_file_recent(path, value[CONF_REFRESH]):
|
||||
_LOGGER.debug("download_gfont: path=%s", path)
|
||||
try:
|
||||
ensure_happy_eyeballs()
|
||||
req = requests.get(url, timeout=external_files.NETWORK_TIMEOUT)
|
||||
req.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
|
||||
@@ -29,6 +29,7 @@ from esphome.const import (
|
||||
UNIT_WATT,
|
||||
)
|
||||
from esphome.core import CORE, HexInt
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
|
||||
DOMAIN = "shelly_dimmer"
|
||||
AUTO_LOAD = ["sensor"]
|
||||
@@ -81,6 +82,7 @@ def get_firmware(value):
|
||||
|
||||
def dl(url):
|
||||
try:
|
||||
ensure_happy_eyeballs()
|
||||
req = requests.get(url, timeout=30)
|
||||
req.raise_for_status()
|
||||
except requests.exceptions.RequestException as e:
|
||||
|
||||
@@ -14,6 +14,7 @@ import requests
|
||||
import esphome.config_validation as cv
|
||||
from esphome.const import CONF_FILE, CONF_TYPE, CONF_URL, __version__
|
||||
from esphome.core import CORE, EsphomeError, TimePeriodSeconds
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
from esphome.helpers import write_file
|
||||
from esphome.types import ConfigType
|
||||
|
||||
@@ -92,6 +93,7 @@ def _write_etag(local_file_path: Path, etag: str | None) -> None:
|
||||
def has_remote_file_changed(
|
||||
url: str, local_file_path: Path, timeout: int = NETWORK_TIMEOUT
|
||||
) -> bool:
|
||||
ensure_happy_eyeballs()
|
||||
if local_file_path.exists():
|
||||
_LOGGER.debug("has_remote_file_changed: File exists at %s", local_file_path)
|
||||
try:
|
||||
@@ -158,6 +160,7 @@ def compute_local_file_dir(domain: str) -> Path:
|
||||
|
||||
|
||||
def download_content(url: str, path: Path, timeout: int = NETWORK_TIMEOUT) -> bytes:
|
||||
ensure_happy_eyeballs()
|
||||
if CORE.skip_external_update and path.exists():
|
||||
_LOGGER.debug("Skipping update for %s (refresh disabled)", url)
|
||||
return path.read_bytes()
|
||||
@@ -231,6 +234,7 @@ def download_content_many(
|
||||
seen: dict[Path, str] = {path: url for url, path in items}
|
||||
if not seen:
|
||||
return
|
||||
ensure_happy_eyeballs()
|
||||
_LOGGER.info("Checking %d %s for updates", len(seen), description)
|
||||
if len(seen) == 1:
|
||||
path, url = next(iter(seen.items()))
|
||||
|
||||
@@ -13,6 +13,7 @@ import sys
|
||||
import time
|
||||
from typing import IO, TYPE_CHECKING
|
||||
|
||||
from esphome.happy_eyeballs import ensure_happy_eyeballs
|
||||
from esphome.helpers import ProgressBar, rmtree
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -755,6 +756,8 @@ def download_with_resume(
|
||||
|
||||
from esphome.core import EsphomeError
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
|
||||
dest = Path(dest)
|
||||
part = dest.with_name(dest.name + ".part")
|
||||
meta = part.with_name(part.name + ".meta")
|
||||
@@ -922,6 +925,8 @@ def download_from_mirrors(
|
||||
|
||||
from esphome.core import EsphomeError
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
|
||||
# 1. Classify the target: filesystem path or open file object
|
||||
path_target: Path | None = None
|
||||
f: IO[bytes] | None = None
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Happy Eyeballs (RFC 8305) connection support for requests/urllib3.
|
||||
|
||||
urllib3 tries each resolved address in sequence with the full connect
|
||||
timeout, so a network advertising IPv6 DNS without IPv6 connectivity stalls
|
||||
every download for the whole timeout before IPv4 is tried.
|
||||
``ensure_happy_eyeballs()`` swaps urllib3's ``create_connection`` for one
|
||||
that races address families with a short stagger via aiohappyeyeballs, run
|
||||
on a daemon-thread event loop so callers stay synchronous.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import socket
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
|
||||
# RFC 8305 recommended delay between staggered connection attempts.
|
||||
HAPPY_EYEBALLS_DELAY = 0.25
|
||||
|
||||
# Extra seconds the connect thread gets beyond the connect timeout before
|
||||
# the caller gives up waiting for it.
|
||||
_THREAD_WAIT_BUFFER = 5.0
|
||||
|
||||
|
||||
def ensure_happy_eyeballs() -> None:
|
||||
"""Make urllib3 (and therefore requests) connect with Happy Eyeballs.
|
||||
|
||||
Idempotent; call before performing requests-based downloads.
|
||||
"""
|
||||
stock: Callable[..., socket.socket] | None = None
|
||||
try:
|
||||
import urllib3.util.connection
|
||||
|
||||
stock = urllib3.util.connection.create_connection
|
||||
if getattr(stock, "_esphome_patched", False):
|
||||
return
|
||||
|
||||
urllib3.util.connection.create_connection = _make_create_connection()
|
||||
except (ImportError, AttributeError) as err: # urllib3 internals moved
|
||||
# WARNING: degraded mode brings back the stalls this module prevents.
|
||||
_LOGGER.warning(
|
||||
"Happy Eyeballs unavailable (%s); downloads use the slower stock "
|
||||
"urllib3 connect",
|
||||
err,
|
||||
)
|
||||
_LOGGER.debug("Happy Eyeballs fallback traceback", exc_info=True)
|
||||
if stock is not None:
|
||||
# Latch so the warning fires once, not per download.
|
||||
stock._esphome_patched = True # type: ignore[attr-defined] # pylint: disable=protected-access
|
||||
|
||||
|
||||
def _make_create_connection() -> Callable[..., socket.socket]:
|
||||
"""Build a drop-in replacement for urllib3's ``create_connection``."""
|
||||
# Deferred so runs that never download skip the ~30 ms asyncio import.
|
||||
import asyncio
|
||||
|
||||
from aiohappyeyeballs import start_connection
|
||||
from urllib3.exceptions import LocationParseError
|
||||
from urllib3.util.connection import ( # noqa: PLC2701
|
||||
_set_socket_options,
|
||||
allowed_gai_family,
|
||||
)
|
||||
from urllib3.util.timeout import _DEFAULT_TIMEOUT # noqa: PLC2701
|
||||
|
||||
from esphome import async_thread
|
||||
|
||||
def create_connection(
|
||||
address: tuple[str, int],
|
||||
timeout: Any = _DEFAULT_TIMEOUT,
|
||||
source_address: tuple[str, int] | None = None,
|
||||
socket_options: Any = None,
|
||||
) -> socket.socket:
|
||||
host, port = address
|
||||
if host.startswith("["):
|
||||
host = host.strip("[]")
|
||||
try:
|
||||
host.encode("idna")
|
||||
except UnicodeError:
|
||||
raise LocationParseError(f"'{host}', label empty or too long") from None
|
||||
|
||||
addr_infos = socket.getaddrinfo(
|
||||
host, port, allowed_gai_family(), socket.SOCK_STREAM
|
||||
)
|
||||
if not addr_infos:
|
||||
# Same error as stock urllib3.
|
||||
raise OSError("getaddrinfo returns an empty list")
|
||||
connect_timeout = (
|
||||
socket.getdefaulttimeout() if timeout is _DEFAULT_TIMEOUT else timeout
|
||||
)
|
||||
|
||||
def socket_factory(addr_info: Any) -> socket.socket:
|
||||
family, type_, proto, _, _ = addr_info
|
||||
sock = socket.socket(family, type_, proto)
|
||||
try:
|
||||
_set_socket_options(sock, socket_options)
|
||||
if source_address:
|
||||
sock.bind(source_address)
|
||||
except BaseException:
|
||||
sock.close()
|
||||
raise
|
||||
return sock
|
||||
|
||||
async def connect() -> socket.socket:
|
||||
return await asyncio.wait_for(
|
||||
start_connection(
|
||||
addr_infos,
|
||||
happy_eyeballs_delay=HAPPY_EYEBALLS_DELAY,
|
||||
interleave=1,
|
||||
socket_factory=socket_factory,
|
||||
),
|
||||
connect_timeout,
|
||||
)
|
||||
|
||||
wait = (
|
||||
None if connect_timeout is None else connect_timeout + _THREAD_WAIT_BUFFER
|
||||
)
|
||||
# on_orphan closes a socket won after the timeout so it cannot leak.
|
||||
sock = async_thread.run_async(
|
||||
connect, timeout=wait, on_orphan=socket.socket.close
|
||||
)
|
||||
# aiohappyeyeballs leaves the winning socket non-blocking; restore the
|
||||
# blocking-with-timeout behavior urllib3 callers expect.
|
||||
try:
|
||||
sock.settimeout(connect_timeout)
|
||||
except BaseException:
|
||||
sock.close()
|
||||
raise
|
||||
return sock
|
||||
|
||||
create_connection._esphome_patched = True # type: ignore[attr-defined] # pylint: disable=protected-access
|
||||
return create_connection
|
||||
@@ -13,6 +13,7 @@ platformio==6.1.19
|
||||
esptool==5.3.1
|
||||
click==8.3.3
|
||||
aioesphomeapi==45.7.0
|
||||
aiohappyeyeballs==2.6.2 # Happy Eyeballs for requests downloads; already pulled in by aioesphomeapi
|
||||
zeroconf==0.150.0
|
||||
puremagic==2.2.0
|
||||
ruamel.yaml==0.19.1 # dashboard_import
|
||||
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Tests for the Happy Eyeballs urllib3 shim."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Generator
|
||||
import socket
|
||||
from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from esphome.happy_eyeballs import _make_create_connection, ensure_happy_eyeballs
|
||||
|
||||
|
||||
def _addr_info(host: str, port: int) -> tuple[Any, ...]:
|
||||
"""Build a getaddrinfo-style result tuple for an IPv4 address."""
|
||||
return (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", (host, port))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def create_connection() -> Any:
|
||||
"""A freshly built Happy Eyeballs create_connection replacement."""
|
||||
return _make_create_connection()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def listener() -> Generator[tuple[str, int]]:
|
||||
"""A listening TCP socket on localhost; yields its address."""
|
||||
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
server.bind(("127.0.0.1", 0))
|
||||
server.listen(5)
|
||||
yield server.getsockname()
|
||||
server.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_gai(listener: tuple[str, int]) -> Generator[Any]:
|
||||
"""Resolve every host to two copies of the listener's address."""
|
||||
with patch("socket.getaddrinfo", return_value=[_addr_info(*listener)] * 2) as mock:
|
||||
yield mock
|
||||
|
||||
|
||||
def test_ensure_happy_eyeballs_patches_and_is_idempotent(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The shim replaces urllib3's create_connection exactly once."""
|
||||
import urllib3.util.connection
|
||||
|
||||
def stock(*args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(urllib3.util.connection, "create_connection", stock)
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
patched = urllib3.util.connection.create_connection
|
||||
assert patched is not stock
|
||||
assert patched._esphome_patched
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
assert urllib3.util.connection.create_connection is patched
|
||||
|
||||
|
||||
def test_connects_and_restores_socket_state(
|
||||
create_connection: Any, listener: tuple[str, int], mock_gai: Any
|
||||
) -> None:
|
||||
"""The winning socket comes back blocking, with timeout and options set."""
|
||||
sock = create_connection(
|
||||
("example.com", listener[1]),
|
||||
timeout=5,
|
||||
socket_options=[(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)],
|
||||
)
|
||||
|
||||
try:
|
||||
assert sock.getpeername() == listener
|
||||
assert sock.gettimeout() == 5
|
||||
assert sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY) != 0
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_single_address_connects(
|
||||
create_connection: Any, listener: tuple[str, int]
|
||||
) -> None:
|
||||
"""A host resolving to one address connects through the same path."""
|
||||
with patch("socket.getaddrinfo", return_value=[_addr_info(*listener)]):
|
||||
sock = create_connection(("example.com", listener[1]), timeout=5)
|
||||
|
||||
try:
|
||||
assert sock.getpeername() == listener
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_falls_back_to_working_address(
|
||||
create_connection: Any, listener: tuple[str, int], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""An unreachable first address does not block the working one."""
|
||||
from esphome import happy_eyeballs
|
||||
|
||||
# 192.0.2.1 (TEST-NET-1) blackholes or fails fast depending on the
|
||||
# network; either way the second address must win well within the
|
||||
# timeout instead of waiting out the first. A short stagger keeps the
|
||||
# test's duration network independent.
|
||||
monkeypatch.setattr(happy_eyeballs, "HAPPY_EYEBALLS_DELAY", 0.01)
|
||||
addr_infos = [_addr_info("192.0.2.1", 9), _addr_info(*listener)]
|
||||
|
||||
with patch("socket.getaddrinfo", return_value=addr_infos):
|
||||
sock = create_connection(("example.com", listener[1]), timeout=10)
|
||||
|
||||
try:
|
||||
assert sock.getpeername() == listener
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_bracketed_ipv6_host_is_stripped(
|
||||
create_connection: Any, listener: tuple[str, int], mock_gai: Any
|
||||
) -> None:
|
||||
"""A bracketed IPv6 literal is unbracketed before resolution."""
|
||||
sock = create_connection(("[::1]", listener[1]), timeout=5)
|
||||
|
||||
try:
|
||||
assert mock_gai.call_args[0][0] == "::1"
|
||||
assert sock.getpeername() == listener
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_source_address_is_bound(
|
||||
create_connection: Any, listener: tuple[str, int], mock_gai: Any
|
||||
) -> None:
|
||||
"""The socket binds to the requested source address before connecting."""
|
||||
sock = create_connection(
|
||||
("example.com", listener[1]),
|
||||
timeout=5,
|
||||
source_address=("127.0.0.1", 0),
|
||||
)
|
||||
|
||||
try:
|
||||
assert sock.getsockname()[0] == "127.0.0.1"
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_socket_factory_failure_closes_socket(
|
||||
listener: tuple[str, int], mock_gai: Any
|
||||
) -> None:
|
||||
"""A socket-option failure fails the connect instead of leaking sockets.
|
||||
|
||||
Instrumented at ``_set_socket_options`` (which the factory calls with
|
||||
the just-created socket) rather than by patching ``socket.socket``,
|
||||
which is platform dependent: the event loop's internal socketpair use
|
||||
differs between platforms.
|
||||
"""
|
||||
created: list[socket.socket] = []
|
||||
|
||||
def failing_set_options(sock: socket.socket, options: Any) -> None:
|
||||
created.append(sock)
|
||||
raise OSError("bad socket option")
|
||||
|
||||
# Patch before building the closure; it binds _set_socket_options at
|
||||
# creation time.
|
||||
with patch("urllib3.util.connection._set_socket_options", new=failing_set_options):
|
||||
create_connection = _make_create_connection()
|
||||
with pytest.raises(OSError):
|
||||
create_connection(
|
||||
("example.com", listener[1]),
|
||||
timeout=5,
|
||||
socket_options=[(999999, 999999, 1)],
|
||||
)
|
||||
|
||||
assert created, "socket factory never ran"
|
||||
assert all(sock.fileno() == -1 for sock in created), "socket leaked open"
|
||||
|
||||
|
||||
def test_default_timeout_yields_blocking_socket(
|
||||
create_connection: Any, listener: tuple[str, int], mock_gai: Any
|
||||
) -> None:
|
||||
"""Without an explicit timeout the socket follows the global default."""
|
||||
sock = create_connection(("example.com", listener[1]))
|
||||
|
||||
try:
|
||||
assert sock.gettimeout() is socket.getdefaulttimeout()
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def test_settimeout_failure_closes_socket(
|
||||
create_connection: Any, mock_gai: Any
|
||||
) -> None:
|
||||
"""A failure restoring socket state closes the winner instead of leaking."""
|
||||
bad_sock = Mock()
|
||||
bad_sock.settimeout.side_effect = OSError("bad timeout")
|
||||
|
||||
with (
|
||||
patch("esphome.async_thread.run_async", return_value=bad_sock),
|
||||
pytest.raises(OSError, match="bad timeout"),
|
||||
):
|
||||
create_connection(("example.com", 80), timeout=5)
|
||||
|
||||
bad_sock.close.assert_called_once()
|
||||
|
||||
|
||||
def test_connect_timeout_raises() -> None:
|
||||
"""A connect that never completes raises within the timeout."""
|
||||
|
||||
async def never(*args: Any, **kwargs: Any) -> None:
|
||||
await asyncio.sleep(60)
|
||||
|
||||
addr_infos = [_addr_info("192.0.2.1", 9), _addr_info("192.0.2.2", 9)]
|
||||
|
||||
# Patch before building the closure; it binds start_connection at
|
||||
# creation time.
|
||||
with patch("aiohappyeyeballs.start_connection", new=never):
|
||||
create_connection = _make_create_connection()
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=addr_infos),
|
||||
pytest.raises(TimeoutError),
|
||||
):
|
||||
create_connection(("example.com", 80), timeout=0.1)
|
||||
|
||||
|
||||
def test_invalid_host_raises_location_parse_error(create_connection: Any) -> None:
|
||||
"""Hostnames urllib3 would reject are still rejected."""
|
||||
from urllib3.exceptions import LocationParseError
|
||||
|
||||
with pytest.raises(LocationParseError):
|
||||
create_connection(("a" * 300, 80))
|
||||
|
||||
|
||||
def test_empty_getaddrinfo_raises_oserror(create_connection: Any) -> None:
|
||||
"""An empty resolution matches stock urllib3's OSError, not ValueError."""
|
||||
with (
|
||||
patch("socket.getaddrinfo", return_value=[]),
|
||||
pytest.raises(OSError, match="empty"),
|
||||
):
|
||||
create_connection(("example.com", 80), timeout=5)
|
||||
|
||||
|
||||
def test_ensure_falls_back_to_stock_when_internals_move(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""If urllib3 private names disappear, downloads keep the stock connect
|
||||
and the warning is latched to fire once, not per download."""
|
||||
import urllib3.util.connection
|
||||
|
||||
from esphome import happy_eyeballs
|
||||
|
||||
def stock(*args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
factory = Mock(side_effect=ImportError("gone"))
|
||||
monkeypatch.setattr(urllib3.util.connection, "create_connection", stock)
|
||||
monkeypatch.setattr(happy_eyeballs, "_make_create_connection", factory)
|
||||
|
||||
ensure_happy_eyeballs()
|
||||
ensure_happy_eyeballs()
|
||||
assert urllib3.util.connection.create_connection is stock
|
||||
assert factory.call_count == 1
|
||||
assert caplog.text.count("Happy Eyeballs unavailable") == 1
|
||||
|
||||
|
||||
def test_ensure_survives_missing_urllib3(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""An unimportable urllib3 degrades with a warning instead of raising."""
|
||||
import sys
|
||||
|
||||
with patch.dict(sys.modules, {"urllib3.util.connection": None}):
|
||||
ensure_happy_eyeballs()
|
||||
assert "Happy Eyeballs unavailable" in caplog.text
|
||||
|
||||
|
||||
def test_requests_routes_through_shim(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Patching urllib3's create_connection actually reroutes requests."""
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
import threading
|
||||
|
||||
import requests
|
||||
import urllib3.util.connection
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", "2")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"ok")
|
||||
|
||||
def log_message(self, *args: Any) -> None:
|
||||
pass
|
||||
|
||||
server = HTTPServer(("127.0.0.1", 0), Handler)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
host, port = server.server_address
|
||||
|
||||
calls: list[Any] = []
|
||||
shim = _make_create_connection()
|
||||
|
||||
def counting(*args: Any, **kwargs: Any) -> Any:
|
||||
calls.append(args)
|
||||
return shim(*args, **kwargs)
|
||||
|
||||
counting._esphome_patched = True
|
||||
monkeypatch.setattr(urllib3.util.connection, "create_connection", counting)
|
||||
|
||||
real_getaddrinfo = socket.getaddrinfo
|
||||
|
||||
def fake_getaddrinfo(h: str, p: int, *args: Any, **kwargs: Any) -> Any:
|
||||
if h == "shim-test.invalid":
|
||||
return [_addr_info(host, port), _addr_info(host, port)]
|
||||
return real_getaddrinfo(h, p, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
|
||||
|
||||
try:
|
||||
with requests.Session() as session:
|
||||
session.trust_env = False
|
||||
resp = session.get(f"http://shim-test.invalid:{port}/", timeout=5)
|
||||
assert resp.status_code == 200
|
||||
assert resp.content == b"ok"
|
||||
assert calls, "requests did not go through the patched create_connection"
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
Reference in New Issue
Block a user