diff --git a/aiopnsense/vnstat.py b/aiopnsense/vnstat.py index 4e17b93..c7d2d31 100644 --- a/aiopnsense/vnstat.py +++ b/aiopnsense/vnstat.py @@ -127,18 +127,26 @@ async def get_vnstat(self) -> MutableMapping[str, Any]: rows_hourly = self._interface_rows(hourly, interface) rows_daily = self._interface_rows(daily, interface) rows_monthly = self._interface_rows(monthly, interface) + parsed_days = [self._parse_daily_label(row.get("label")) for row in rows_daily] + parsed_months = [self._parse_month_label(row.get("label")) for row in rows_monthly] selected_rows = { "vnstat_today": self._pick_daily_row( - rows_daily, days_ago=0, current_tz=opnsense_tz + rows_daily, days_ago=0, current_tz=opnsense_tz, parsed_days=parsed_days ), "vnstat_this_month": self._pick_monthly_row( - rows_monthly, months_ago=0, current_tz=opnsense_tz + rows_monthly, + months_ago=0, + current_tz=opnsense_tz, + parsed_months=parsed_months, ), "vnstat_yesterday": self._pick_daily_row( - rows_daily, days_ago=1, current_tz=opnsense_tz + rows_daily, days_ago=1, current_tz=opnsense_tz, parsed_days=parsed_days ), "vnstat_last_month": self._pick_monthly_row( - rows_monthly, months_ago=1, current_tz=opnsense_tz + rows_monthly, + months_ago=1, + current_tz=opnsense_tz, + parsed_months=parsed_months, ), "vnstat_last_hour": self._pick_last_hour_row(rows_hourly, current_tz=opnsense_tz), } @@ -304,7 +312,11 @@ def _to_bits_per_second(self, value: str, unit: str) -> int | None: return int(round(parsed_value * factor)) def _pick_daily_row( - self, rows: Sequence[dict[str, Any]], days_ago: int, current_tz: tzinfo + self, + rows: Sequence[dict[str, Any]], + days_ago: int, + current_tz: tzinfo, + parsed_days: Sequence[date | None] | None = None, ) -> dict[str, Any] | None: """Select a daily row by matching day label or falling back by position. @@ -312,6 +324,8 @@ def _pick_daily_row( rows (Sequence[dict[str, Any]]): Collection of parsed table rows. days_ago (int): Day offset used for fallback selection. current_tz (tzinfo): Timezone used to determine the current date. + parsed_days (Sequence[date | None] | None): Pre-parsed row labels, + when available for reuse across selections. Returns: dict[str, Any] | None: Daily row matching ``days_ago`` in @@ -319,8 +333,8 @@ def _pick_daily_row( second-latest row for yesterday when labels cannot be parsed. """ target_day = datetime.now(tz=current_tz).date() - timedelta(days=days_ago) - for row in rows: - parsed_day = self._parse_daily_label(row.get("label")) + days = parsed_days or [self._parse_daily_label(row.get("label")) for row in rows] + for row, parsed_day in zip(rows, days, strict=False): if parsed_day == target_day: return row if days_ago == 0 and rows: @@ -330,7 +344,11 @@ def _pick_daily_row( return None def _pick_monthly_row( - self, rows: Sequence[dict[str, Any]], months_ago: int, current_tz: tzinfo + self, + rows: Sequence[dict[str, Any]], + months_ago: int, + current_tz: tzinfo, + parsed_months: Sequence[tuple[int, int] | None] | None = None, ) -> dict[str, Any] | None: """Select a monthly row by matching month label or fallback position. @@ -338,6 +356,9 @@ def _pick_monthly_row( rows (Sequence[dict[str, Any]]): Collection of parsed table rows. months_ago (int): Month offset used for fallback selection. current_tz (tzinfo): Timezone used to determine the current month. + parsed_months (Sequence[tuple[int, int] | None] | None): + Pre-parsed row labels, when available for reuse across + selections. Returns: dict[str, Any] | None: Monthly row matching ``months_ago`` in @@ -352,8 +373,8 @@ def _pick_monthly_row( target_month += 12 target_year -= 1 - for row in rows: - parsed_month = self._parse_month_label(row.get("label")) + months = parsed_months or [self._parse_month_label(row.get("label")) for row in rows] + for row, parsed_month in zip(rows, months, strict=False): if parsed_month == (target_year, target_month): return row if months_ago == 0 and rows: diff --git a/tests/test_vnstat.py b/tests/test_vnstat.py index b8827e0..a1a4cf9 100644 --- a/tests/test_vnstat.py +++ b/tests/test_vnstat.py @@ -3,7 +3,7 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest @@ -213,6 +213,8 @@ async def fake_safe_get(path: str, *_args: Any, **_kwargs: Any) -> dict[str, Any return {} client._safe_dict_get = AsyncMock(side_effect=fake_safe_get) + client._parse_daily_label = Mock(wraps=client._parse_daily_label) + client._parse_month_label = Mock(wraps=client._parse_month_label) vnstat = await client.get_vnstat() gib = 1024**3 @@ -234,6 +236,8 @@ async def fake_safe_get(path: str, *_args: Any, **_kwargs: Any) -> dict[str, Any assert igc1_metrics["vnstat_yesterday"]["total_bytes"] == 1 * gib assert igc1_metrics["vnstat_last_month"]["total_bytes"] == 2 * gib assert igc1_metrics["vnstat_last_hour"]["total_bytes"] == int(1.5 * gib) + assert client._parse_daily_label.call_count == 4 + assert client._parse_month_label.call_count == 4 finally: await client.async_close() @@ -341,12 +345,19 @@ async def test_parse_vnstat_payload_and_helpers_edge_cases(make_client: ClientTy assert client._pick_daily_row([{"label": "bad0"}, {"label": "bad1"}], 1, tz) == { "label": "bad0" } + assert client._pick_daily_row([{"label": "bad"}], 2, tz) is None this_month = datetime.now(tz=tz).strftime("%b '%y") assert client._pick_monthly_row([{"label": this_month}], 0, tz) == {"label": this_month} assert client._pick_monthly_row([{"label": "bad0"}, {"label": "bad1"}], 1, tz) == { "label": "bad0" } + now = datetime.now(tz=tz) + previous_year = now.year - 1 + previous_december = f"Dec '{previous_year % 100:02d}" + assert client._pick_monthly_row([{"label": previous_december}], now.month, tz) == { + "label": previous_december + } now_hour = datetime.now(tz=tz).replace(minute=0, second=0, microsecond=0) prev_hour = now_hour - timedelta(hours=1)