From 034fabe188c183e93f689c2a88790b93125ddd78 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 6 Sep 2023 04:04:49 -0500
Subject: [PATCH] Use loop time to set context (#99701)

* Use loop time to set context

loop time is faster than utcnow, and since its only used internally it can
be switched without a breaking change

* fix mocking
---
 homeassistant/helpers/entity.py | 11 ++++++-----
 tests/helpers/test_service.py   |  7 ++++---
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py
index e946c41d3b8..7bd510b6fa1 100644
--- a/homeassistant/helpers/entity.py
+++ b/homeassistant/helpers/entity.py
@@ -5,7 +5,7 @@ from abc import ABC
 import asyncio
 from collections.abc import Coroutine, Iterable, Mapping, MutableMapping
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import timedelta
 from enum import Enum, auto
 import functools as ft
 import logging
@@ -41,7 +41,7 @@ from homeassistant.exceptions import (
     NoEntitySpecifiedError,
 )
 from homeassistant.loader import bind_hass
-from homeassistant.util import dt as dt_util, ensure_unique_string, slugify
+from homeassistant.util import ensure_unique_string, slugify
 
 from . import device_registry as dr, entity_registry as er
 from .device_registry import DeviceInfo, EventDeviceRegistryUpdatedData
@@ -272,7 +272,7 @@ class Entity(ABC):
 
     # Context
     _context: Context | None = None
-    _context_set: datetime | None = None
+    _context_set: float | None = None
 
     # If entity is added to an entity platform
     _platform_state = EntityPlatformState.NOT_ADDED
@@ -660,7 +660,7 @@ class Entity(ABC):
     def async_set_context(self, context: Context) -> None:
         """Set the context the entity currently operates under."""
         self._context = context
-        self._context_set = dt_util.utcnow()
+        self._context_set = self.hass.loop.time()
 
     async def async_update_ha_state(self, force_refresh: bool = False) -> None:
         """Update Home Assistant with current state of entity.
@@ -847,7 +847,8 @@ class Entity(ABC):
 
         if (
             self._context_set is not None
-            and dt_util.utcnow() - self._context_set > self.context_recent_time
+            and hass.loop.time() - self._context_set
+            > self.context_recent_time.total_seconds()
         ):
             self._context = None
             self._context_set = None
diff --git a/tests/helpers/test_service.py b/tests/helpers/test_service.py
index 803a57e12ed..03a8b5e11b2 100644
--- a/tests/helpers/test_service.py
+++ b/tests/helpers/test_service.py
@@ -1,5 +1,4 @@
 """Test service helpers."""
-from collections import OrderedDict
 from collections.abc import Iterable
 from copy import deepcopy
 from typing import Any
@@ -54,7 +53,7 @@ def mock_handle_entity_call():
 
 
 @pytest.fixture
-def mock_entities(hass):
+def mock_entities(hass: HomeAssistant) -> dict[str, MockEntity]:
     """Return mock entities in an ordered dict."""
     kitchen = MockEntity(
         entity_id="light.kitchen",
@@ -80,11 +79,13 @@ def mock_entities(hass):
         should_poll=False,
         supported_features=(SUPPORT_B | SUPPORT_C),
     )
-    entities = OrderedDict()
+    entities = {}
     entities[kitchen.entity_id] = kitchen
     entities[living_room.entity_id] = living_room
     entities[bedroom.entity_id] = bedroom
     entities[bathroom.entity_id] = bathroom
+    for entity in entities.values():
+        entity.hass = hass
     return entities