Merge dc07077abd into 8c99ac78b7

src/netsleuth/probes/icmp_win.py

@@ -48,9 +48,12 @@ def parse_echo_reply(buffer: bytes, pointer_size: int = 8) -> IcmpReply:
     ttl = buffer[ttl_offset]
     if status in (IP_REQ_TIMED_OUT,) or (a, b, c, d) == (0, 0, 0, 0):
         return IcmpReply(address=None, status=status, rtt_ms=None, ttl=ttl)
-    return IcmpReply(address=f"{a}.{b}.{c}.{d}", status=status, rtt_ms=float(rtt), ttl=ttl)
+    return IcmpReply(
+        address=f"{a}.{b}.{c}.{d}", status=status, rtt_ms=float(rtt), ttl=ttl
+    )
 
 
+import asyncio
 import ctypes
 import platform
 import socket
@@ -139,7 +142,11 @@ def echo_once(
     iphlpapi, handle = _handle()
     try:
         options = _IpOptionInformation(
-            Ttl=ttl, Tos=0, Flags=IP_FLAG_DF if df else 0, OptionsSize=0, OptionsData=None
+            Ttl=ttl,
+            Tos=0,
+            Flags=IP_FLAG_DF if df else 0,
+            OptionsSize=0,
+            OptionsData=None,
         )
         buffer = ctypes.create_string_buffer(_REPLY_BUFFER_SIZE)
         address = struct.unpack("<I", socket.inet_aton(dest))[0]
@@ -180,7 +187,7 @@ def echo_once(
         iphlpapi.IcmpCloseHandle(handle)
 
 
-def ping_samples_win(
+async def ping_samples_win(
     host: str, count: int, interval: float, timeout: float, source_ip: str | None = None
 ) -> list[float | None]:
     dest = socket.gethostbyname(host)
@@ -188,8 +195,10 @@ def ping_samples_win(
     samples: list[float | None] = []
     for index in range(count):
         if index:
-            time.sleep(interval)
+            await asyncio.sleep(interval)
-        reply = echo_once(dest, ttl=128, timeout_ms=timeout_ms, source_ip=source_ip)
+        reply = await asyncio.to_thread(
+            echo_once, dest, ttl=128, timeout_ms=timeout_ms, source_ip=source_ip
+        )
         samples.append(reply.rtt_ms if classify_status(reply.status) == "ok" else None)
     return samples
 
@@ -207,12 +216,16 @@ def trace_hops_win(
     # into 30 minutes of blocking IcmpSendEcho2 calls.
     address = socket.gethostbyname(dest)
     hops: list[tuple[int, IcmpReply]] = []
-    per_hop_timeout_ms = max(1, min(timeout_ms // max(max_hops, 1), _MAX_HOP_TIMEOUT_MS))
+    per_hop_timeout_ms = max(
+        1, min(timeout_ms // max(max_hops, 1), _MAX_HOP_TIMEOUT_MS)
+    )
     deadline = time.monotonic() + timeout_ms / 1000.0
     for ttl in range(1, max_hops + 1):
         if time.monotonic() >= deadline:
             break
-        reply = echo_once(address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip)
+        reply = echo_once(
+            address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip
+        )
         hops.append((ttl, reply))
         if classify_status(reply.status) == "ok":
             break

src/netsleuth/probes/latency.py

@@ -47,7 +47,7 @@ async def tcp_connect_rtt(
         _reader, writer = await asyncio.wait_for(
             asyncio.open_connection(host, port, local_addr=local_addr), timeout=timeout
         )
-    except (TimeoutError, OSError):
+    except TimeoutError, OSError:
         return None
     writer.close()
     try:
@@ -62,22 +62,32 @@ async def _resolve(host: str) -> str | None:
         loop = asyncio.get_running_loop()
         res = await loop.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
         return res[0][4][0]
-    except (OSError, IndexError):
+    except OSError, IndexError:
         return None
 
 
 async def _icmp_samples(
-    host: str, count: int, interval: float, timeout: float, backend: str, source_ip: str | None = None
+    host: str,
+    count: int,
+    interval: float,
+    timeout: float,
+    backend: str,
+    source_ip: str | None = None,
 ) -> list[float | None]:
     if backend == "icmp_win":
         from netsleuth.probes.icmp_win import ping_samples_win
 
-        return await asyncio.to_thread(ping_samples_win, host, count, interval, timeout, source_ip)
+        return await ping_samples_win(host, count, interval, timeout, source_ip)
     from icmplib import async_ping
 
     privileged = backend == "icmp_raw"
     host_result = await async_ping(
-        host, count=count, interval=interval, timeout=timeout, privileged=privileged, source=source_ip
+        host,
+        count=count,
+        interval=interval,
+        timeout=timeout,
+        privileged=privileged,
+        source=source_ip,
     )
     samples: list[float | None] = list(host_result.rtts)
     samples.extend([None] * (count - len(samples)))
@@ -96,11 +106,15 @@ async def ping_host(
 ) -> PingResult:
     if semaphore is not None:
         async with semaphore:
-            return await ping_host(host, label, count, interval, timeout, backend, source_ip)
+            return await ping_host(
+                host, label, count, interval, timeout, backend, source_ip
+            )
     resolved = await _resolve(host)
     if backend in ("icmp_win", "icmp_dgram", "icmp_raw"):
         try:
-            samples = await _icmp_samples(host, count, interval, timeout, backend, source_ip)
+            samples = await _icmp_samples(
+                host, count, interval, timeout, backend, source_ip
+            )
             return summarize_ping(label, host, resolved, backend, samples)
         except Exception:
             backend = "tcp"
@@ -108,7 +122,9 @@ async def ping_host(
     for index in range(count):
         if index:
             await asyncio.sleep(interval)
-        samples.append(await tcp_connect_rtt(host, timeout=timeout, source_ip=source_ip))
+        samples.append(
+            await tcp_connect_rtt(host, timeout=timeout, source_ip=source_ip)
+        )
     return summarize_ping(label, host, resolved, "tcp", samples)
 
 
@@ -125,7 +141,9 @@ async def ping_fanout(
     return list(
         await asyncio.gather(
             *(
-                ping_host(host, label, count, interval, timeout, backend, source_ip, semaphore)
+                ping_host(
+                    host, label, count, interval, timeout, backend, source_ip, semaphore
+                )
                 for label, host in hosts
             )
         )

tests/test_latency.py

@@ -6,7 +6,12 @@ from unittest.mock import AsyncMock, patch, MagicMock
 import pytest
 
 from netsleuth.models import PingResult
-from netsleuth.probes.latency import summarize_ping, tcp_connect_rtt, _resolve, _icmp_samples
+from netsleuth.probes.latency import (
+    summarize_ping,
+    tcp_connect_rtt,
+    _resolve,
+    _icmp_samples,
+)
 
 
 def test_summarize_ping_calculates_correct_stats():
@@ -68,6 +73,7 @@ async def test_tcp_connect_rtt_oserror(mock_open_conn):
     rtt = await tcp_connect_rtt("host.com")
     assert rtt is None
 
+
 @pytest.mark.asyncio
 async def test_resolve_success():
     # Mock loop.getaddrinfo via asyncio.get_running_loop
@@ -75,9 +81,7 @@ async def test_resolve_success():
     with patch.object(loop, "getaddrinfo", new_callable=AsyncMock) as mock_getaddrinfo:
         # getaddrinfo returns list of tuples: (family, type, proto, canonname, sockaddr)
         # sockaddr is (address, port) for IPv4
-        mock_getaddrinfo.return_value = [
+        mock_getaddrinfo.return_value = [(2, 1, 6, "", ("1.2.3.4", 0))]
-            (2, 1, 6, '', ('1.2.3.4', 0))
-        ]
         ip = await _resolve("host.com")
         assert ip == "1.2.3.4"
 
@@ -92,17 +96,12 @@ async def test_resolve_failure():
 
 
 @pytest.mark.asyncio
-@patch("netsleuth.probes.latency.asyncio.to_thread", new_callable=AsyncMock)
+@patch("netsleuth.probes.icmp_win.ping_samples_win", new_callable=AsyncMock)
-async def test_icmp_samples_win(mock_to_thread):
+async def test_icmp_samples_win(mock_ping_samples_win):
-    mock_to_thread.return_value = [10.0, 15.0]
+    mock_ping_samples_win.return_value = [10.0, 15.0]
     samples = await _icmp_samples("host.com", 2, 0.5, 2.0, "icmp_win", source_ip=None)
     assert samples == [10.0, 15.0]
-    mock_to_thread.assert_called_once()
+    mock_ping_samples_win.assert_called_once_with("host.com", 2, 0.5, 2.0, None)
-    # The first argument should be the ping_samples_win function
-    func_arg = mock_to_thread.call_args[0][0]
-    assert func_arg.__name__ == "ping_samples_win"
-    # The remaining arguments
-    assert mock_to_thread.call_args[0][1:] == ("host.com", 2, 0.5, 2.0, None)
 
 
 @pytest.mark.asyncio
@@ -133,7 +132,12 @@ async def test_icmp_samples_unix_dgram_padded():
         # Should be padded with Nones to length 3
         assert samples == [10.0, None, None]
         mock_async_ping.assert_called_once_with(
-            "host.com", count=3, interval=0.5, timeout=2.0, privileged=False, source=None
+            "host.com",
+            count=3,
+            interval=0.5,
+            timeout=2.0,
+            privileged=False,
+            source=None,
         )
 
 
@@ -160,7 +164,9 @@ async def test_ping_host_icmp_success(mock_icmp_samples, mock_resolve):
 @patch("netsleuth.probes.latency._resolve", new_callable=AsyncMock)
 @patch("netsleuth.probes.latency._icmp_samples", new_callable=AsyncMock)
 @patch("netsleuth.probes.latency.tcp_connect_rtt", new_callable=AsyncMock)
-async def test_ping_host_icmp_exception_fallback(mock_tcp_connect, mock_icmp_samples, mock_resolve):
+async def test_ping_host_icmp_exception_fallback(
+    mock_tcp_connect, mock_icmp_samples, mock_resolve
+):
     mock_resolve.return_value = "1.2.3.4"
     mock_icmp_samples.side_effect = Exception("ICMP failed")
     mock_tcp_connect.side_effect = [15.0, 16.0]
@@ -190,7 +196,9 @@ async def test_ping_host_tcp_direct(mock_tcp_connect, mock_resolve):
 async def test_ping_host_semaphore():
     semaphore = asyncio.Semaphore(1)
 
-    with patch("netsleuth.probes.latency.ping_host", new_callable=AsyncMock) as mock_ping_host:
+    with patch(
+        "netsleuth.probes.latency.ping_host", new_callable=AsyncMock
+    ) as mock_ping_host:
         mock_ping_host.return_value = MagicMock()
 
         # When called with a semaphore, ping_host calls itself without the semaphore inside the block