Fix tests for 2023.11

dext0r · Oct 28, 2023 · 076f872 · 076f872
1 parent b5bd48f
commit 076f872
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 12 deletions.
diff --git a/tests/test_cloud.py b/tests/test_cloud.py
@@ -6,9 +6,9 @@
 from aiohttp import WSMessage, WSMsgType
 from homeassistant import core
 from homeassistant.components import demo
-from homeassistant.const import Platform
+from homeassistant.const import MAJOR_VERSION, MINOR_VERSION, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import DATA_CLIENTSESSION
+from homeassistant.helpers.aiohttp_client import DATA_CLIENTSESSION, _make_key
 from homeassistant.setup import async_setup_component
 import pytest
 from pytest_homeassistant_custom_component.common import MockConfigEntry
@@ -59,14 +59,20 @@ async def ws_connect(self, *args, **kwargs):
         return self.ws
 
 
+def mock_client_session(hass, session):
+    if MAJOR_VERSION > 2023 or (MAJOR_VERSION == 2023 and MINOR_VERSION >= 11):
+        hass.data[DATA_CLIENTSESSION] = {_make_key(): session}
+    else:
+        hass.data[DATA_CLIENTSESSION] = session
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: MockConfigEntry,
     session: MockSession | None = None,
     aiohttp_client: Any | None = None,
 ):
-    hass.data[DATA_CLIENTSESSION] = session or MockSession(aiohttp_client)
-
+    mock_client_session(hass, session or MockSession(aiohttp_client))
     config_entry.add_to_hass(hass)
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()

diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
@@ -1,7 +1,6 @@
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.config import YAML_CONFIG_FILE
 from homeassistant.const import CONF_ENTITIES
-from homeassistant.helpers.aiohttp_client import DATA_CLIENTSESSION
 from homeassistant.helpers.reload import async_integration_yaml_config
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.setup import async_setup_component
@@ -65,7 +64,7 @@ async def test_config_flow_empty_entities(hass):
 
 
 async def test_config_flow_cloud(hass, aioclient_mock):
-    hass.data[DATA_CLIENTSESSION] = test_cloud.MockSession(aioclient_mock)
+    test_cloud.mock_client_session(hass, test_cloud.MockSession(aioclient_mock))
 
     result = await hass.config_entries.flow.async_init(DOMAIN, context={"source": config_entries.SOURCE_USER})
     assert result["type"] == "form"

diff --git a/tests/test_notifier.py b/tests/test_notifier.py
@@ -9,7 +9,6 @@
 from homeassistant.components.light import ATTR_BRIGHTNESS
 from homeassistant.const import ATTR_DEVICE_CLASS, EVENT_HOMEASSISTANT_STARTED, STATE_UNAVAILABLE
 from homeassistant.core import CoreState, State
-from homeassistant.helpers.aiohttp_client import DATA_CLIENTSESSION
 from homeassistant.helpers.template import Template
 from homeassistant.setup import async_setup_component
 import pytest
@@ -82,14 +81,14 @@ async def test_notifier_setup_config_invalid(hass, hass_admin_user, config_entry
     component: YandexSmartHome = hass.data[DOMAIN]
     with pytest.raises(KeyError):
         component.get_entry_data(config_entry_direct)
-    assert (
-        caplog.messages[-1] == "Config entry 'Mock Title' for yandex_smart_home integration not ready yet: "
-        "User invalid does not exist; Retrying in background"
+    assert caplog.messages[-1].startswith(
+        "Config entry 'Mock Title' for yandex_smart_home integration not ready yet: "
+        "User invalid does not exist; Retrying in"
     )
 
 
 async def test_notifier_setup_not_discovered(hass, hass_admin_user, aioclient_mock):
-    hass.data[DATA_CLIENTSESSION] = test_cloud.MockSession(aioclient_mock)
+    test_cloud.mock_client_session(hass, test_cloud.MockSession(aioclient_mock))
 
     yaml_config = {
         const.CONF_NOTIFIER: [
@@ -129,7 +128,7 @@ async def test_notifier_setup_not_discovered(hass, hass_admin_user, aioclient_mo
 
 
 async def test_notifier_lifecycle_discovered(hass, hass_admin_user, aioclient_mock):
-    hass.data[DATA_CLIENTSESSION] = test_cloud.MockSession(aioclient_mock)
+    test_cloud.mock_client_session(hass, test_cloud.MockSession(aioclient_mock))
 
     yaml_config = {
         const.CONF_NOTIFIER: [