Fix timezone edge cases for Unifi Protect media source (#77636)

* Fixes timezone edge cases for Unifi Protect media source

* linting

homeassistant/components/unifiprotect/media_source.py

@@ -101,12 +101,12 @@ async def async_get_media_source(hass: HomeAssistant) -> MediaSource:
 
 
 @callback
-def _get_start_end(hass: HomeAssistant, start: datetime) -> tuple[datetime, datetime]:
+def _get_month_start_end(start: datetime) -> tuple[datetime, datetime]:
     start = dt_util.as_local(start)
     end = dt_util.now()
 
-    start = start.replace(day=1, hour=1, minute=0, second=0, microsecond=0)
+    start = start.replace(day=1, hour=0, minute=0, second=1, microsecond=0)
-    end = end.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+    end = end.replace(day=1, hour=0, minute=0, second=2, microsecond=0)
 
     return start, end
 
@@ -571,9 +571,16 @@ class ProtectMediaSource(MediaSource):
         if not build_children:
             return source
 
-        month = start.month
+        if data.api.bootstrap.recording_start is not None:
+            recording_start = data.api.bootstrap.recording_start.date()
+        start = max(recording_start, start)
+
+        recording_end = dt_util.now().date()
+        end = start.replace(month=start.month + 1) - timedelta(days=1)
+        end = min(recording_end, end)
+
         children = [self._build_days(data, camera_id, event_type, start, is_all=True)]
-        while start.month == month:
+        while start <= end:
             children.append(
                 self._build_days(data, camera_id, event_type, start, is_all=False)
             )
@@ -702,7 +709,7 @@ class ProtectMediaSource(MediaSource):
             self._build_recent(data, camera_id, event_type, 30),
         ]
 
-        start, end = _get_start_end(self.hass, data.api.bootstrap.recording_start)
+        start, end = _get_month_start_end(data.api.bootstrap.recording_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

@@ -4,7 +4,9 @@ from datetime import datetime, timedelta
 from ipaddress import IPv4Address
 from unittest.mock import AsyncMock, Mock, patch
 
+from freezegun import freeze_time
 import pytest
+import pytz
 from pyunifiprotect.data import (
     Bootstrap,
     Camera,
@@ -28,6 +30,7 @@ from homeassistant.components.unifiprotect.media_source import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.util import dt as dt_util
 
 from .conftest import MockUFPFixture
 from .utils import init_entry
@@ -430,13 +433,52 @@ async def test_browse_media_event_type(
     assert browse.children[3].identifier == "test_id:browse:all:smart"
 
 
+ONE_MONTH_SIMPLE = (
+    datetime(
+        year=2022,
+        month=9,
+        day=1,
+        hour=3,
+        minute=0,
+        second=0,
+        microsecond=0,
+        tzinfo=pytz.timezone("US/Pacific"),
+    ),
+    1,
+)
+TWO_MONTH_SIMPLE = (
+    datetime(
+        year=2022,
+        month=8,
+        day=31,
+        hour=3,
+        minute=0,
+        second=0,
+        microsecond=0,
+        tzinfo=pytz.timezone("US/Pacific"),
+    ),
+    2,
+)
+
+
+@pytest.mark.parametrize(
+    "start,months",
+    [ONE_MONTH_SIMPLE, TWO_MONTH_SIMPLE],
+)
+@freeze_time("2022-09-15 03:00:00-07:00")
 async def test_browse_media_time(
-    hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, fixed_now: datetime
+    hass: HomeAssistant,
+    ufp: MockUFPFixture,
+    doorbell: Camera,
+    start: datetime,
+    months: int,
 ):
     """Test browsing time selector level media."""
 
-    last_month = fixed_now.replace(day=1) - timedelta(days=1)
+    end = datetime.fromisoformat("2022-09-15 03:00:00-07:00")
-    ufp.api.bootstrap._recording_start = last_month
+    end_local = dt_util.as_local(end)
+
+    ufp.api.bootstrap._recording_start = dt_util.as_utc(start)
 
     ufp.api.get_bootstrap = AsyncMock(return_value=ufp.api.bootstrap)
     await init_entry(hass, ufp, [doorbell], regenerate_ids=False)
@@ -449,17 +491,89 @@ async def test_browse_media_time(
 
     assert browse.title == f"UnifiProtect > {doorbell.name} > All Events"
     assert browse.identifier == base_id
-    assert len(browse.children) == 4
+    assert len(browse.children) == 3 + months
     assert browse.children[0].title == "Last 24 Hours"
     assert browse.children[0].identifier == f"{base_id}:recent:1"
     assert browse.children[1].title == "Last 7 Days"
     assert browse.children[1].identifier == f"{base_id}:recent:7"
     assert browse.children[2].title == "Last 30 Days"
     assert browse.children[2].identifier == f"{base_id}:recent:30"
-    assert browse.children[3].title == f"{fixed_now.strftime('%B %Y')}"
+    assert browse.children[3].title == f"{end_local.strftime('%B %Y')}"
     assert (
         browse.children[3].identifier
-        == f"{base_id}:range:{fixed_now.year}:{fixed_now.month}"
+        == f"{base_id}:range:{end_local.year}:{end_local.month}"
+    )
+
+
+ONE_MONTH_TIMEZONE = (
+    datetime(
+        year=2022,
+        month=8,
+        day=1,
+        hour=3,
+        minute=0,
+        second=0,
+        microsecond=0,
+        tzinfo=pytz.timezone("US/Pacific"),
+    ),
+    1,
+)
+TWO_MONTH_TIMEZONE = (
+    datetime(
+        year=2022,
+        month=7,
+        day=31,
+        hour=21,
+        minute=0,
+        second=0,
+        microsecond=0,
+        tzinfo=pytz.timezone("US/Pacific"),
+    ),
+    2,
+)
+
+
+@pytest.mark.parametrize(
+    "start,months",
+    [ONE_MONTH_TIMEZONE, TWO_MONTH_TIMEZONE],
+)
+@freeze_time("2022-08-31 21:00:00-07:00")
+async def test_browse_media_time_timezone(
+    hass: HomeAssistant,
+    ufp: MockUFPFixture,
+    doorbell: Camera,
+    start: datetime,
+    months: int,
+):
+    """Test browsing time selector level media."""
+
+    end = datetime.fromisoformat("2022-08-31 21:00:00-07:00")
+    end_local = dt_util.as_local(end)
+
+    ufp.api.bootstrap._recording_start = dt_util.as_utc(start)
+
+    ufp.api.get_bootstrap = AsyncMock(return_value=ufp.api.bootstrap)
+    await init_entry(hass, ufp, [doorbell], regenerate_ids=False)
+
+    base_id = f"test_id:browse:{doorbell.id}:all"
+    source = await async_get_media_source(hass)
+    media_item = MediaSourceItem(hass, DOMAIN, base_id, None)
+
+    browse = await source.async_browse_media(media_item)
+
+    assert browse.title == f"UnifiProtect > {doorbell.name} > All Events"
+    assert browse.identifier == base_id
+    assert len(browse.children) == 3 + months
+    assert browse.children[0].title == "Last 24 Hours"
+    assert browse.children[0].identifier == f"{base_id}:recent:1"
+    assert browse.children[1].title == "Last 7 Days"
+    assert browse.children[1].identifier == f"{base_id}:recent:7"
+    assert browse.children[2].title == "Last 30 Days"
+    assert browse.children[2].identifier == f"{base_id}:recent:30"
+    assert browse.children[3].title == f"{end_local.strftime('%B %Y')}"
+    assert (
+        browse.children[3].identifier
+        == f"{base_id}:range:{end_local.year}:{end_local.month}"
     )
 
 
@@ -599,13 +713,14 @@ async def test_browse_media_eventthumb(
     assert browse.media_class == MEDIA_CLASS_IMAGE
 
 
+@freeze_time("2022-09-15 03:00:00-07:00")
 async def test_browse_media_day(
     hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, fixed_now: datetime
 ):
     """Test browsing day selector level media."""
 
-    last_month = fixed_now.replace(day=1) - timedelta(days=1)
+    start = datetime.fromisoformat("2022-09-03 03:00:00-07:00")
-    ufp.api.bootstrap._recording_start = last_month
+    ufp.api.bootstrap._recording_start = dt_util.as_utc(start)
 
     ufp.api.get_bootstrap = AsyncMock(return_value=ufp.api.bootstrap)
     await init_entry(hass, ufp, [doorbell], regenerate_ids=False)
@@ -623,7 +738,7 @@ async def test_browse_media_day(
         == f"UnifiProtect > {doorbell.name} > All Events > {fixed_now.strftime('%B %Y')}"
     )
     assert browse.identifier == base_id
-    assert len(browse.children) in (29, 30, 31, 32)
+    assert len(browse.children) == 14
     assert browse.children[0].title == "Whole Month"
     assert browse.children[0].identifier == f"{base_id}:all"