From c191a7cfdb0c6cde0825d14683a041edb10f9442 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Fri, 4 Oct 2024 08:24:01 +0200
Subject: [PATCH] Fix lingering tasks in snooz tests (#127523)

---
 tests/components/snooz/__init__.py  | 33 ++++++++++++++++++++++++++++-
 tests/components/snooz/test_fan.py  |  7 +++---
 tests/components/snooz/test_init.py |  6 ------
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/tests/components/snooz/__init__.py b/tests/components/snooz/__init__.py
index 16827b54ac4..f27ef91fe5a 100644
--- a/tests/components/snooz/__init__.py
+++ b/tests/components/snooz/__init__.py
@@ -6,7 +6,8 @@ from dataclasses import dataclass
 from unittest.mock import patch
 
 from pysnooz.commands import SnoozCommandData
-from pysnooz.testing import MockSnoozDevice
+from pysnooz.device import DisconnectionReason, SnoozConnectionStatus
+from pysnooz.testing import MockSnoozDevice as ParentMockSnoozDevice
 
 from homeassistant.components.snooz.const import DOMAIN
 from homeassistant.const import CONF_ADDRESS, CONF_TOKEN
@@ -66,6 +67,36 @@ class SnoozFixture:
     device: MockSnoozDevice
 
 
+class MockSnoozDevice(ParentMockSnoozDevice):
+    """Used for testing integration with Bleak.
+
+    Adjusted for https://github.com/AustinBrunkhorst/pysnooz/pull/19
+    """
+
+    async def async_disconnect(self) -> None:
+        """Disconnect from the device."""
+        self._is_manually_disconnecting = True
+        try:
+            self._cancel_current_command()
+            if (
+                self._reconnection_task is not None
+                and not self._reconnection_task.done()
+            ):
+                self._reconnection_task.cancel()
+
+            if self._connection_task is not None and not self._connection_task.done():
+                self._connection_task.cancel()
+
+            if self._api is not None:
+                await self._api.async_disconnect()
+
+            if self.connection_status != SnoozConnectionStatus.DISCONNECTED:
+                self._machine.device_disconnected(reason=DisconnectionReason.USER)
+
+        finally:
+            self._is_manually_disconnecting = False
+
+
 async def create_mock_snooz(
     connected: bool = True,
     initial_state: SnoozCommandData = SnoozCommandData(on=False, volume=0),
diff --git a/tests/components/snooz/test_fan.py b/tests/components/snooz/test_fan.py
index ddc93a4ba1f..69b06692557 100644
--- a/tests/components/snooz/test_fan.py
+++ b/tests/components/snooz/test_fan.py
@@ -149,8 +149,6 @@ async def test_transition_off(hass: HomeAssistant, snooz_fan_entity_id: str) ->
     assert ATTR_ASSUMED_STATE not in state.attributes
 
 
-# This tests needs to be adjusted to remove lingering tasks
-@pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_push_events(
     hass: HomeAssistant, mock_connected_snooz: SnoozFixture, snooz_fan_entity_id: str
 ) -> None:
@@ -174,9 +172,10 @@ async def test_push_events(
     state = hass.states.get(snooz_fan_entity_id)
     assert state.attributes[ATTR_ASSUMED_STATE] is True
 
+    # Don't attempt to reconnect
+    await mock_connected_snooz.device.async_disconnect()
+
 
-# This tests needs to be adjusted to remove lingering tasks
-@pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_restore_state(
     hass: HomeAssistant, entity_registry: er.EntityRegistry
 ) -> None:
diff --git a/tests/components/snooz/test_init.py b/tests/components/snooz/test_init.py
index b1ab06fcc8e..edcd7913792 100644
--- a/tests/components/snooz/test_init.py
+++ b/tests/components/snooz/test_init.py
@@ -2,15 +2,11 @@
 
 from __future__ import annotations
 
-import pytest
-
 from homeassistant.core import HomeAssistant
 
 from . import SnoozFixture
 
 
-# This tests needs to be adjusted to remove lingering tasks
-@pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_removing_entry_cleans_up_connections(
     hass: HomeAssistant, mock_connected_snooz: SnoozFixture
 ) -> None:
@@ -21,8 +17,6 @@ async def test_removing_entry_cleans_up_connections(
     assert not mock_connected_snooz.device.is_connected
 
 
-# This tests needs to be adjusted to remove lingering tasks
-@pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_reloading_entry_cleans_up_connections(
     hass: HomeAssistant, mock_connected_snooz: SnoozFixture
 ) -> None: