From f11aba96486743ca4e8ab40c4d430b840d649a05 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 4 Nov 2024 00:25:37 +0100
Subject: [PATCH] Fix flaky tests in advantage_air (#129758)

---
 .../advantage_air/test_binary_sensor.py       | 44 ++++++-------------
 tests/components/advantage_air/test_sensor.py | 24 +++-------
 2 files changed, 20 insertions(+), 48 deletions(-)

diff --git a/tests/components/advantage_air/test_binary_sensor.py b/tests/components/advantage_air/test_binary_sensor.py
index 7a7b2f8df5b..d0088d96ba5 100644
--- a/tests/components/advantage_air/test_binary_sensor.py
+++ b/tests/components/advantage_air/test_binary_sensor.py
@@ -1,10 +1,8 @@
 """Test the Advantage Air Binary Sensor Platform."""
 
 from datetime import timedelta
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
-from homeassistant.components.advantage_air import ADVANTAGE_AIR_SYNC_INTERVAL
-from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
@@ -70,22 +68,14 @@ async def test_binary_sensor_async_setup_entry(
     assert not hass.states.get(entity_id)
 
     mock_get.reset_mock()
-    entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
-    await hass.async_block_till_done()
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=ADVANTAGE_AIR_SYNC_INTERVAL + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 1
+    with patch("homeassistant.config_entries.RELOAD_AFTER_UPDATE_DELAY", 1):
+        entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
+        await hass.async_block_till_done()
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 3
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
+        await hass.async_block_till_done(wait_background_tasks=True)
+        assert len(mock_get.mock_calls) == 1
 
     state = hass.states.get(entity_id)
     assert state
@@ -101,22 +91,14 @@ async def test_binary_sensor_async_setup_entry(
     assert not hass.states.get(entity_id)
 
     mock_get.reset_mock()
-    entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
-    await hass.async_block_till_done()
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=ADVANTAGE_AIR_SYNC_INTERVAL + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 1
+    with patch("homeassistant.config_entries.RELOAD_AFTER_UPDATE_DELAY", 1):
+        entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
+        await hass.async_block_till_done()
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 3
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
+        await hass.async_block_till_done(wait_background_tasks=True)
+        assert len(mock_get.mock_calls) == 1
 
     state = hass.states.get(entity_id)
     assert state
diff --git a/tests/components/advantage_air/test_sensor.py b/tests/components/advantage_air/test_sensor.py
index 4389e67228a..3ea368a59fb 100644
--- a/tests/components/advantage_air/test_sensor.py
+++ b/tests/components/advantage_air/test_sensor.py
@@ -1,15 +1,13 @@
 """Test the Advantage Air Sensor Platform."""
 
 from datetime import timedelta
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
-from homeassistant.components.advantage_air import ADVANTAGE_AIR_SYNC_INTERVAL
 from homeassistant.components.advantage_air.const import DOMAIN as ADVANTAGE_AIR_DOMAIN
 from homeassistant.components.advantage_air.sensor import (
     ADVANTAGE_AIR_SERVICE_SET_TIME_TO,
     ADVANTAGE_AIR_SET_COUNTDOWN_VALUE,
 )
-from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
@@ -124,23 +122,15 @@ async def test_sensor_platform_disabled_entity(
 
     assert not hass.states.get(entity_id)
 
-    entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
-    await hass.async_block_till_done(wait_background_tasks=True)
     mock_get.reset_mock()
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=ADVANTAGE_AIR_SYNC_INTERVAL + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 1
+    with patch("homeassistant.config_entries.RELOAD_AFTER_UPDATE_DELAY", 1):
+        entity_registry.async_update_entity(entity_id=entity_id, disabled_by=None)
+        await hass.async_block_till_done(wait_background_tasks=True)
 
-    async_fire_time_changed(
-        hass,
-        dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1),
-    )
-    await hass.async_block_till_done(wait_background_tasks=True)
-    assert len(mock_get.mock_calls) == 3
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
+        await hass.async_block_till_done(wait_background_tasks=True)
+        assert len(mock_get.mock_calls) == 1
 
     state = hass.states.get(entity_id)
     assert state