diff --git a/src/netsleuth/probes/icmp_win.py b/src/netsleuth/probes/icmp_win.py index 1b242f5..f90d542 100644 --- a/src/netsleuth/probes/icmp_win.py +++ b/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(" 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) - reply = echo_once(dest, ttl=128, timeout_ms=timeout_ms, source_ip=source_ip) + await asyncio.sleep(interval) + 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 diff --git a/src/netsleuth/probes/latency.py b/src/netsleuth/probes/latency.py index 1927d24..eeeb467 100644 --- a/src/netsleuth/probes/latency.py +++ b/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 ) ) diff --git a/tests/test_latency.py b/tests/test_latency.py index d51dcfa..7a384f5 100644 --- a/tests/test_latency.py +++ b/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 = [ - (2, 1, 6, '', ('1.2.3.4', 0)) - ] + mock_getaddrinfo.return_value = [(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) -async def test_icmp_samples_win(mock_to_thread): - mock_to_thread.return_value = [10.0, 15.0] +@patch("netsleuth.probes.icmp_win.ping_samples_win", new_callable=AsyncMock) +async def test_icmp_samples_win(mock_ping_samples_win): + 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() - # 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) + mock_ping_samples_win.assert_called_once_with("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