Skip to content

Commit

Permalink
Update modified_at datetime on storage collection changes (home-assis…
Browse files Browse the repository at this point in the history
  • Loading branch information
edenhaus authored and bramkragten committed Sep 4, 2024
1 parent de99dfe commit 122f11c
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 3 deletions.
37 changes: 34 additions & 3 deletions homeassistant/helpers/collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from collections.abc import Awaitable, Callable, Coroutine, Iterable
from dataclasses import dataclass
from functools import partial
from hashlib import md5
from itertools import groupby
import logging
from operator import attrgetter
Expand All @@ -25,6 +26,7 @@
from . import entity_registry
from .entity import Entity
from .entity_component import EntityComponent
from .json import json_bytes
from .storage import Store
from .typing import ConfigType, VolDictType

Expand All @@ -50,6 +52,7 @@ class CollectionChange:
change_type: str
item_id: str
item: Any
item_hash: str | None = None


type ChangeListener = Callable[
Expand Down Expand Up @@ -273,7 +276,9 @@ async def async_load(self) -> None:
await self.notify_changes(
[
CollectionChange(CHANGE_ADDED, item[CONF_ID], item)
CollectionChange(
CHANGE_ADDED, item[CONF_ID], item, self._hash_item(item)
)
for item in raw_storage["items"]
]
)
Expand Down Expand Up @@ -313,7 +318,16 @@ async def async_create_item(self, data: dict) -> _ItemT:
item = self._create_item(item_id, validated_data)
self.data[item_id] = item
self._async_schedule_save()
await self.notify_changes([CollectionChange(CHANGE_ADDED, item_id, item)])
await self.notify_changes(
[
CollectionChange(
CHANGE_ADDED,
item_id,
item,
self._hash_item(self._serialize_item(item_id, item)),
)
]
)
return item

async def async_update_item(self, item_id: str, updates: dict) -> _ItemT:
Expand All @@ -331,7 +345,16 @@ async def async_update_item(self, item_id: str, updates: dict) -> _ItemT:
self.data[item_id] = updated
self._async_schedule_save()

await self.notify_changes([CollectionChange(CHANGE_UPDATED, item_id, updated)])
await self.notify_changes(
[
CollectionChange(
CHANGE_UPDATED,
item_id,
updated,
self._hash_item(self._serialize_item(item_id, updated)),
)
]
)

return self.data[item_id]

Expand Down Expand Up @@ -365,6 +388,10 @@ def _base_data_to_save(self) -> SerializedStorageCollection:
def _data_to_save(self) -> _StoreT:
"""Return JSON-compatible date for storing to file."""

def _hash_item(self, item: dict) -> str:
"""Return a hash of the item."""
return md5(json_bytes(item)).hexdigest()


class DictStorageCollection(StorageCollection[dict, SerializedStorageCollection]):
"""A specialized StorageCollection where the items are untyped dicts."""
Expand Down Expand Up @@ -464,6 +491,10 @@ async def _remove_entity(self, change_set: CollectionChange) -> None:

async def _update_entity(self, change_set: CollectionChange) -> None:
if entity := self.entities.get(change_set.item_id):
if change_set.item_hash:
self.ent_reg.async_update_entity_options(
entity.entity_id, "collection", {"hash": change_set.item_hash}
)
await entity.async_update_config(change_set.item)

async def _collection_changed(self, change_set: Iterable[CollectionChange]) -> None:
Expand Down
81 changes: 81 additions & 0 deletions tests/helpers/test_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@

from __future__ import annotations

from datetime import timedelta
import logging

from freezegun.api import FrozenDateTimeFactory
import pytest
import voluptuous as vol

Expand All @@ -15,6 +17,7 @@
storage,
)
from homeassistant.helpers.typing import ConfigType
from homeassistant.util.dt import utcnow

from tests.common import flush_store
from tests.typing import WebSocketGenerator
Expand Down Expand Up @@ -254,6 +257,84 @@ async def test_storage_collection(hass: HomeAssistant) -> None:
}


async def test_storage_collection_update_modifiet_at(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
freezer: FrozenDateTimeFactory,
) -> None:
"""Test that updating a storage collection will update the modified_at datetime in the entity registry."""

entities: dict[str, TestEntity] = {}

class TestEntity(MockEntity):
"""Entity that is config based."""

def __init__(self, config: ConfigType) -> None:
"""Initialize entity."""
super().__init__(config)
self._state = "initial"

@classmethod
def from_storage(cls, config: ConfigType) -> TestEntity:
"""Create instance from storage."""
obj = super().from_storage(config)
entities[obj.unique_id] = obj
return obj

@property
def state(self) -> str:
"""Return state of entity."""
return self._state

def set_state(self, value: str) -> None:
"""Set value."""
self._state = value
self.async_write_ha_state()

store = storage.Store(hass, 1, "test-data")
data = {"id": "mock-1", "name": "Mock 1", "data": 1}
await store.async_save(
{
"items": [
data,
]
}
)
id_manager = collection.IDManager()
ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
await ent_comp.async_setup({})
coll = MockStorageCollection(store, id_manager)
collection.sync_entity_lifecycle(hass, "test", "test", ent_comp, coll, TestEntity)
changes = track_changes(coll)

await coll.async_load()
assert id_manager.has_id("mock-1")
assert len(changes) == 1
assert changes[0] == (collection.CHANGE_ADDED, "mock-1", data)

modified_1 = entity_registry.async_get("test.mock_1").modified_at
assert modified_1 == utcnow()

freezer.tick(timedelta(minutes=1))

updated_item = await coll.async_update_item("mock-1", {"data": 2})
assert id_manager.has_id("mock-1")
assert updated_item == {"id": "mock-1", "name": "Mock 1", "data": 2}
assert len(changes) == 2
assert changes[1] == (collection.CHANGE_UPDATED, "mock-1", updated_item)

modified_2 = entity_registry.async_get("test.mock_1").modified_at
assert modified_2 > modified_1
assert modified_2 == utcnow()

freezer.tick(timedelta(minutes=1))

entities["mock-1"].set_state("second")

modified_3 = entity_registry.async_get("test.mock_1").modified_at
assert modified_3 == modified_2


async def test_attach_entity_component_collection(hass: HomeAssistant) -> None:
"""Test attaching collection to entity component."""
ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
Expand Down

0 comments on commit 122f11c

Please sign in to comment.