Improve date handling in UniFi Protect media source (#159491)

Co-authored-by: RaHehl <rahehl@users.noreply.github.com>

homeassistant/components/unifiprotect/media_source.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from calendar import monthrange
 from datetime import date, datetime, timedelta
 from enum import Enum
 from typing import Any, NoReturn, cast
@@ -94,11 +95,12 @@ async def async_get_media_source(hass: HomeAssistant) -> MediaSource:
 
 @callback
 def _get_month_start_end(start: datetime) -> tuple[datetime, datetime]:
+    """Get the first day of the month for start and current time."""
     start = dt_util.as_local(start)
     end = dt_util.now()
 
-    start = start.replace(day=1, hour=0, minute=0, second=1, microsecond=0)
-    end = end.replace(day=1, hour=0, minute=0, second=2, microsecond=0)
+    start = start.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+    end = end.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
 
     return start, end
 
@@ -113,20 +115,19 @@ def _bad_identifier(identifier: str, err: Exception | None = None) -> NoReturn:
 
 @callback
 def _format_duration(duration: timedelta) -> str:
-    formatted = ""
     seconds = int(duration.total_seconds())
-    if seconds > 3600:
-        hours = seconds // 3600
-        formatted += f"{hours}h "
-        seconds -= hours * 3600
-    if seconds > 60:
-        minutes = seconds // 60
-        formatted += f"{minutes}m "
-        seconds -= minutes * 60
-    if seconds > 0:
-        formatted += f"{seconds}s "
+    hours, seconds = divmod(seconds, 3600)
+    minutes, seconds = divmod(seconds, 60)
 
-    return formatted.strip()
+    parts = []
+    if hours > 0:
+        parts.append(f"{hours}h")
+    if minutes > 0:
+        parts.append(f"{minutes}m")
+    if seconds > 0:
+        parts.append(f"{seconds}s")
+
+    return " ".join(parts) if parts else "0s"
 
 
 @callback
@@ -593,7 +594,8 @@ class ProtectMediaSource(MediaSource):
         start = max(recording_start, start)
 
         recording_end = dt_util.now().date()
-        end = start.replace(month=start.month + 1) - timedelta(days=1)
+
+        end = start.replace(day=monthrange(start.year, start.month)[1])
         end = min(recording_end, end)
 
         children = [self._build_days(data, camera_id, event_type, start, is_all=True)]
@@ -660,10 +662,9 @@ class ProtectMediaSource(MediaSource):
             tzinfo=dt_util.get_default_time_zone(),
         )
         if is_all:
-            if start_dt.month < 12:
-                end_dt = start_dt.replace(month=start_dt.month + 1)
-            else:
-                end_dt = start_dt.replace(year=start_dt.year + 1, month=1)
+            # Move to first day of next month
+            days_in_month = monthrange(start_dt.year, start_dt.month)[1]
+            end_dt = start_dt + timedelta(days=days_in_month)
         else:
             end_dt = start_dt + timedelta(hours=24)
 
@@ -726,7 +727,7 @@ class ProtectMediaSource(MediaSource):
         ]
 
         start, end = _get_month_start_end(data.api.bootstrap.recording_start)
-        while end > start:
+        while end >= start:
             children.append(self._build_month(data, camera_id, event_type, end.date()))
             end = (end - timedelta(days=1)).replace(day=1)
 

tests/components/unifiprotect/test_media_source.py

@@ -21,6 +21,7 @@ from homeassistant.components.media_source import MediaSourceItem
 from homeassistant.components.unifiprotect.const import DOMAIN
 from homeassistant.components.unifiprotect.media_source import (
     ProtectMediaSource,
+    SimpleEventType,
     async_get_media_source,
 )
 from homeassistant.core import HomeAssistant
@@ -1041,3 +1042,66 @@ async def test_browse_media_browse_whole_month_december(
     assert browse.identifier == base_id
     assert len(browse.children) == 1
     assert browse.children[0].identifier == "test_id:event:test_event_id"
+
+
+@pytest.mark.parametrize(
+    ("year", "month", "expected_days", "expected_end_month", "expected_end_year"),
+    [
+        (2024, 1, 31, 2, 2024),  # January
+        (2024, 2, 29, 3, 2024),  # February (leap year)
+        (2023, 2, 28, 3, 2023),  # February (non-leap year)
+        (2024, 4, 30, 5, 2024),  # April
+        (2024, 12, 31, 1, 2025),  # December - critical edge case
+    ],
+)
+async def test_build_days_whole_month_date_calculation(
+    hass: HomeAssistant,
+    ufp: MockUFPFixture,
+    year: int,
+    month: int,
+    expected_days: int,
+    expected_end_month: int,
+    expected_end_year: int,
+) -> None:
+    """Test that whole month date calculation works for all month types.
+
+    This test verifies the monthrange-based date calculation in _build_days,
+    especially for December which previously used manual year/month increment logic.
+    """
+    # Initialize the integration entry to get ProtectData
+    await init_entry(hass, ufp, [], regenerate_ids=False)
+
+    # Create a start date for the first day of the month
+    start = datetime(year=year, month=month, day=1).date()
+    start_dt = datetime(
+        year=start.year,
+        month=start.month,
+        day=start.day,
+        hour=0,
+        minute=0,
+        second=0,
+        tzinfo=dt_util.get_default_time_zone(),
+    )
+
+    # Verify we got the expected number of days
+    expected_end = start_dt + timedelta(days=expected_days)
+
+    # Verify it correctly goes to the expected month/year
+    assert expected_end.month == expected_end_month
+    assert expected_end.year == expected_end_year
+    assert expected_end.day == 1
+
+    # Build the media source with is_all=True (whole month)
+    source = ProtectMediaSource(hass, {})
+    result = await source._build_days(
+        data=ufp.entry.runtime_data,
+        camera_id="test_camera",
+        event_type=SimpleEventType.ALL,
+        start=start,
+        is_all=True,
+        build_children=False,  # We only care about the identifier, not children
+    )
+
+    # Verify the identifier format is correct
+    assert result.identifier.endswith(f"range:{year}:{month}:all")
+    assert "Whole Month" in result.title