Repository: XKNX/xknx
Pull request #1489 (2024-06-16 10:01:52)

Non-awaitable callbacks in ConnectionManager
<!--
  You are awesome! Thanks for contributing to our project!
  Please, DO NOT DELETE ANY TEXT from this template!.
-->
## Description
<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.
-->


Fixes # (issue)

## Type of change
<!--
Please tick the applicable options.
NOTE: Ticking multiple options most likely indicates
that your change is to big and it is suggested to split it into several smaller PRs.
-->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [x] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Code quality improvements to existing code or addition of tests

## Checklist

- [x] The documentation has been adjusted accordingly
- [x] Tests have been added that prove the fix is effective or that the feature works
- [x] The changes are documented in the changelog (docs/changelog.md)

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/XKNX/xknx/pull/1489?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX) Report
Attention: Patch coverage is `90.00000%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 96.63%. Comparing base [(`3d5777e`)](https://app.codecov.io/gh/XKNX/xknx/commit/3d5777ed1d453e826b8b1a6281af9f4c1c08e059?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX) to head [(`499395b`)](https://app.codecov.io/gh/XKNX/xknx/commit/499395b2e769cfe9eaf3e5f8d6130246d6f44fc4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX).

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/XKNX/xknx/pull/1489/graphs/tree.svg?width=650&height=150&src=pr&token=irWbIygS84&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX)](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX)

```diff
@@            Coverage Diff             @@
##             main    #1489      +/-   ##
==========================================
- Coverage   96.63%   96.63%   -0.01%     
==========================================
  Files         154      154              
  Lines       10026    10025       -1     
==========================================
- Hits         9689     9688       -1     
  Misses        337      337              
```


| [Files](https://app.codecov.io/gh/XKNX/xknx/pull/1489?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX) | Coverage Δ | |
|---|---|---|
| [xknx/core/connection\_manager.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fcore%2Fconnection_manager.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC9jb3JlL2Nvbm5lY3Rpb25fbWFuYWdlci5weQ==) | `100.00% <100.00%> (ø)` | |
| [xknx/core/state\_updater.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fcore%2Fstate_updater.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC9jb3JlL3N0YXRlX3VwZGF0ZXIucHk=) | `83.47% <100.00%> (ø)` | |
| [xknx/core/task\_registry.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fcore%2Ftask_registry.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC9jb3JlL3Rhc2tfcmVnaXN0cnkucHk=) | `100.00% <100.00%> (ø)` | |
| [xknx/io/tunnel.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fio%2Ftunnel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC9pby90dW5uZWwucHk=) | `85.82% <100.00%> (ø)` | |
| [xknx/typing/\_\_init\_\_.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Ftyping%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC90eXBpbmcvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [xknx/xknx.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fxknx.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC94a254LnB5) | `88.77% <100.00%> (-0.12%)` | :arrow_down: |
| [xknx/io/routing.py](https://app.codecov.io/gh/XKNX/xknx/pull/1489?src=pr&el=tree&filepath=xknx%2Fio%2Frouting.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XKNX#diff-eGtueC9pby9yb3V0aW5nLnB5) | `92.12% <50.00%> (ø)` | |

</details>

Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 38508b34..8231611a 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -11,9 +11,10 @@ nav_order: 2
 ### Breaking changes
 
 - Drop support for Python 3.9
-- Change callback signatures from awaitable to callable in `XKNX.device_updated_cb`, TelegramQueue, Device, Devices and RemoteValue.
+- Change callback signatures from awaitable to callable in `XKNX.device_updated_cb`, TelegramQueue, Device, Devices, ConnectionManager and RemoteValue.
 - Remove `async` from functions / methods (nothing has to be awaited there)
   - Tools:  `group_value_write`, `group_value_response` and `group_value_read`
+  - ConnectionManager: `.connection_state_changed`
   - RemoteValue: `.set`, `.respond`, `.process` and `.update_value`
   - ValueReader: `.send_group_read`
 
diff --git a/docs/xknx.md b/docs/xknx.md
index 40151f74..535d4ea5 100644
--- a/docs/xknx.md
+++ b/docs/xknx.md
@@ -177,7 +177,7 @@ async def main():
 asyncio.run(main())
 ```
 
-An awaitable `connection_state_changed_cb` will be called every time the connection state to the gateway changes. Example:
+A callback `connection_state_changed_cb` will be called every time the connection state to the gateway changes. Example:
 
 ```python
 import asyncio
@@ -185,7 +185,7 @@ from xknx import XKNX
 from xknx.core import XknxConnectionState
 
 
-async def connection_state_changed_cb(state: XknxConnectionState):
+def connection_state_changed_cb(state: XknxConnectionState):
     print("Callback received with state {0}".format(state.name))
 
 
diff --git a/xknx/core/connection_manager.py b/xknx/core/connection_manager.py
index 42404ead..eb561d79 100644
--- a/xknx/core/connection_manager.py
+++ b/xknx/core/connection_manager.py
@@ -3,12 +3,10 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Callable
 from datetime import datetime
 
 from xknx.core.connection_state import XknxConnectionState, XknxConnectionType
-
-AsyncConnectionStateCallback = Callable[[XknxConnectionState], Awaitable[None]]
+from xknx.typing import ConnectionChangeCallbackType
 
 
 class ConnectionManager:
@@ -20,7 +18,7 @@ class ConnectionManager:
 
         self.connected = asyncio.Event()
         self._state = XknxConnectionState.DISCONNECTED
-        self._connection_state_changed_cbs: list[AsyncConnectionStateCallback] = []
+        self._connection_state_changed_cbs: list[ConnectionChangeCallbackType] = []
 
         self.cemi_count_incoming: int = 0
         self.cemi_count_incoming_error: int = 0
@@ -34,32 +32,32 @@ class ConnectionManager:
         self._main_loop = asyncio.get_running_loop()
 
     def register_connection_state_changed_cb(
-        self, connection_state_changed_cb: AsyncConnectionStateCallback
+        self, connection_state_changed_cb: ConnectionChangeCallbackType
     ) -> None:
         """Register callback for connection state being updated."""
         self._connection_state_changed_cbs.append(connection_state_changed_cb)
 
     def unregister_connection_state_changed_cb(
-        self, connection_state_changed_cb: AsyncConnectionStateCallback
+        self, connection_state_changed_cb: ConnectionChangeCallbackType
     ) -> None:
         """Unregister callback for connection state being updated."""
         if connection_state_changed_cb in self._connection_state_changed_cbs:
             self._connection_state_changed_cbs.remove(connection_state_changed_cb)
 
-    async def connection_state_changed(
+    def connection_state_changed(
         self,
         state: XknxConnectionState,
         connection_type: XknxConnectionType = XknxConnectionType.NOT_CONNECTED,
     ) -> None:
         """Run registered callbacks in main loop. Set internal state flag."""
         if self._main_loop:
-            asyncio.run_coroutine_threadsafe(
-                self._connection_state_changed(state, connection_type), self._main_loop
+            self._main_loop.call_soon_threadsafe(
+                self._connection_state_changed, state, connection_type
             )
         else:
-            await self._connection_state_changed(state, connection_type)
+            self._connection_state_changed(state, connection_type)
 
-    async def _connection_state_changed(
+    def _connection_state_changed(
         self, state: XknxConnectionState, connection_type: XknxConnectionType
     ) -> None:
         """Run registered callbacks. Set internal state flag."""
@@ -75,11 +73,8 @@ class ConnectionManager:
             self.connected.clear()
             self.connected_since = None
 
-        if tasks := [
+        for connection_state_change_cb in self._connection_state_changed_cbs:
             connection_state_change_cb(state)
-            for connection_state_change_cb in self._connection_state_changed_cbs
-        ]:
-            await asyncio.gather(*tasks)
 
     @property
     def state(self) -> XknxConnectionState:
diff --git a/xknx/core/state_updater.py b/xknx/core/state_updater.py
index 80cbc783..233d67fd 100644
--- a/xknx/core/state_updater.py
+++ b/xknx/core/state_updater.py
@@ -203,9 +203,7 @@ class StateUpdater:
 
         self._stop()
 
-    async def connection_state_change_callback(
-        self, state: XknxConnectionState
-    ) -> None:
+    def connection_state_change_callback(self, state: XknxConnectionState) -> None:
         """Start and stop StateUpdater via connection state update."""
         if state == XknxConnectionState.CONNECTED:
             if not self.started:
diff --git a/xknx/core/task_registry.py b/xknx/core/task_registry.py
index 39d0626d..9dccd21d 100644
--- a/xknx/core/task_registry.py
+++ b/xknx/core/task_registry.py
@@ -127,7 +127,7 @@ class TaskRegistry:
         """Await all tracked tasks."""
         await asyncio.gather(*self.tasks)
 
-    async def connection_state_changed_cb(self, state: XknxConnectionState) -> None:
+    def connection_state_changed_cb(self, state: XknxConnectionState) -> None:
         """Handle connection state changes."""
         for task in self.tasks:
             if state == XknxConnectionState.CONNECTED:
diff --git a/xknx/io/routing.py b/xknx/io/routing.py
index fec70f18..e23c9b8b 100644
--- a/xknx/io/routing.py
+++ b/xknx/io/routing.py
@@ -181,7 +181,7 @@ class Routing(Interface):
     async def connect(self) -> bool:
         """Start routing."""
         self.xknx.current_address = self.individual_address
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTING, self.connection_type
         )
         try:
@@ -192,13 +192,13 @@ class Routing(Interface):
                 type(ex).__name__,
                 ex,
             )
-            await self.xknx.connection_manager.connection_state_changed(
+            self.xknx.connection_manager.connection_state_changed(
                 XknxConnectionState.DISCONNECTED
             )
             # close udp transport to prevent open file descriptors
             self.transport.stop()
             raise CommunicationError("Routing could not be started") from ex
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTED, self.connection_type
         )
         return True
@@ -206,7 +206,7 @@ class Routing(Interface):
     async def disconnect(self) -> None:
         """Stop routing."""
         self.transport.stop()
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
         self._flow_control.cancel()
diff --git a/xknx/io/tunnel.py b/xknx/io/tunnel.py
index 6a13b954..f36782ae 100644
--- a/xknx/io/tunnel.py
+++ b/xknx/io/tunnel.py
@@ -100,7 +100,7 @@ class _Tunnel(Interface):
 
     async def connect(self) -> bool:
         """Connect to a KNX tunneling interface. Returns True on success."""
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTING, self.connection_type
         )
         try:
@@ -113,7 +113,7 @@ class _Tunnel(Interface):
                 type(ex).__name__,
                 ex,
             )
-            await self.xknx.connection_manager.connection_state_changed(
+            self.xknx.connection_manager.connection_state_changed(
                 XknxConnectionState.DISCONNECTED
             )
             if not self._initial_connection and self.auto_reconnect:
@@ -126,7 +126,7 @@ class _Tunnel(Interface):
             ) from ex
 
         self._tunnel_established()
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTED, self.connection_type
         )
         return True
@@ -140,10 +140,8 @@ class _Tunnel(Interface):
     def _tunnel_lost(self) -> None:
         """Prepare for reconnection or shutdown when the connection is lost. Callback."""
         self.stop_heartbeat()
-        self.xknx.task_registry.background(
-            self.xknx.connection_manager.connection_state_changed(
-                XknxConnectionState.DISCONNECTED
-            )
+        self.xknx.connection_manager.connection_state_changed(
+            XknxConnectionState.DISCONNECTED
         )
         self._data_endpoint_addr = None
         if self.auto_reconnect:
@@ -169,7 +167,7 @@ class _Tunnel(Interface):
     async def disconnect(self) -> None:
         """Disconnect tunneling connection."""
         self.stop_heartbeat()
-        await self.xknx.connection_manager.connection_state_changed(
+        self.xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
         self._data_endpoint_addr = None
diff --git a/xknx/typing/__init__.py b/xknx/typing/__init__.py
index b3b190b8..f3cff2f4 100644
--- a/xknx/typing/__init__.py
+++ b/xknx/typing/__init__.py
@@ -4,11 +4,14 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, TypeVar
 
 if TYPE_CHECKING:
+    from xknx.core.connection_manager import XknxConnectionState
     from xknx.devices import Device
     from xknx.telegram import Telegram
 
 CallbackType = Callable[[], None]
 
+ConnectionChangeCallbackType = Callable[["XknxConnectionState"], None]
+
 DeviceT = TypeVar("DeviceT", bound="Device")
 DeviceCallbackType = Callable[[DeviceT], None]
 
diff --git a/xknx/xknx.py b/xknx/xknx.py
index 138d1d58..4db14622 100644
--- a/xknx/xknx.py
+++ b/xknx/xknx.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Callable
 import logging
 from logging.handlers import TimedRotatingFileHandler
 from pathlib import Path
@@ -17,7 +16,6 @@ from xknx.core import (
     GroupAddressDPT,
     TaskRegistry,
     TelegramQueue,
-    XknxConnectionState,
 )
 from xknx.core.state_updater import StateUpdater, TrackerOptionType
 from xknx.devices import Device, Devices
@@ -29,7 +27,11 @@ from xknx.io import (
 )
 from xknx.management import Management
 from xknx.telegram import GroupAddress, GroupAddressType, IndividualAddress, Telegram
-from xknx.typing import DeviceCallbackType, TelegramCallbackType
+from xknx.typing import (
+    ConnectionChangeCallbackType,
+    DeviceCallbackType,
+    TelegramCallbackType,
+)
 
 from .__version__ import __version__ as VERSION
 
@@ -44,8 +46,7 @@ class XKNX:
         address_format: GroupAddressType = GroupAddressType.LONG,
         telegram_received_cb: TelegramCallbackType | None = None,
         device_updated_cb: DeviceCallbackType[Device] | None = None,
-        connection_state_changed_cb: Callable[[XknxConnectionState], Awaitable[None]]
-        | None = None,
+        connection_state_changed_cb: ConnectionChangeCallbackType | None = None,
         rate_limit: int = 0,
         multicast_group: str = DEFAULT_MCAST_GRP,
         multicast_port: int = DEFAULT_MCAST_PORT,


Test patch:
diff --git a/test/core_tests/connection_manager_test.py b/test/core_tests/connection_manager_test.py
index 5cc1d5ad..c0084686 100644
--- a/test/core_tests/connection_manager_test.py
+++ b/test/core_tests/connection_manager_test.py
@@ -1,8 +1,9 @@
 """Unit test for connection manager."""
 
+import asyncio
 from datetime import datetime
 import threading
-from unittest.mock import AsyncMock, patch
+from unittest.mock import Mock, patch
 
 from xknx import XKNX
 from xknx.core import XknxConnectionState, XknxConnectionType
@@ -20,14 +21,14 @@ class TestConnectionManager:
         """Test connection_state_changed after register."""
 
         xknx = XKNX()
-        async_connection_state_changed_cb = AsyncMock()
+        connection_state_changed_cb = Mock()
         xknx.connection_manager.register_connection_state_changed_cb(
-            async_connection_state_changed_cb
+            connection_state_changed_cb
         )
-        await xknx.connection_manager.connection_state_changed(
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTED, XknxConnectionType.ROUTING_SECURE
         )
-        async_connection_state_changed_cb.assert_called_once_with(
+        connection_state_changed_cb.assert_called_once_with(
             XknxConnectionState.CONNECTED
         )
 
@@ -35,17 +36,15 @@ class TestConnectionManager:
         """Test unregister after register."""
 
         xknx = XKNX()
-        async_connection_state_changed_cb = AsyncMock()
+        connection_state_changed_cb = Mock()
         xknx.connection_manager.register_connection_state_changed_cb(
-            async_connection_state_changed_cb
+            connection_state_changed_cb
         )
         xknx.connection_manager.unregister_connection_state_changed_cb(
-            async_connection_state_changed_cb
+            connection_state_changed_cb
         )
-        await xknx.connection_manager.connection_state_changed(
-            XknxConnectionState.CONNECTED
-        )
-        async_connection_state_changed_cb.assert_not_called()
+        xknx.connection_manager.connection_state_changed(XknxConnectionState.CONNECTED)
+        connection_state_changed_cb.assert_not_called()
 
     #
     # TEST PROCESS
@@ -54,15 +53,15 @@ class TestConnectionManager:
         """Test should return if current state equals parameter."""
 
         xknx = XKNX()
-        async_connection_state_changed_cb = AsyncMock()
+        connection_state_changed_cb = Mock()
         xknx.connection_manager.register_connection_state_changed_cb(
-            async_connection_state_changed_cb
+            connection_state_changed_cb
         )
         assert xknx.connection_manager.state == XknxConnectionState.DISCONNECTED
-        await xknx.connection_manager.connection_state_changed(
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
-        async_connection_state_changed_cb.assert_not_called()
+        connection_state_changed_cb.assert_not_called()
 
     #
     # TEST CONNECTED
@@ -71,17 +70,15 @@ class TestConnectionManager:
         """Test connected event callback."""
 
         xknx = XKNX()
-        async_connection_state_changed_cb = AsyncMock()
+        connection_state_changed_cb = Mock()
         xknx.connection_manager.register_connection_state_changed_cb(
-            async_connection_state_changed_cb
+            connection_state_changed_cb
         )
 
         assert not xknx.connection_manager.connected.is_set()
 
-        await xknx.connection_manager.connection_state_changed(
-            XknxConnectionState.CONNECTED
-        )
-        async_connection_state_changed_cb.assert_called_once_with(
+        xknx.connection_manager.connection_state_changed(XknxConnectionState.CONNECTED)
+        connection_state_changed_cb.assert_called_once_with(
             XknxConnectionState.CONNECTED
         )
 
@@ -93,7 +90,7 @@ class TestConnectionManager:
         self.main_thread = threading.get_ident()
         xknx = XKNX(connection_config=ConnectionConfig(threaded=True))
 
-        async def assert_main_thread(*args, **kwargs):
+        def assert_main_thread(*args, **kwargs):
             """Test callback is done by main thread."""
             assert self.main_thread == threading.get_ident()
 
@@ -101,7 +98,7 @@ class TestConnectionManager:
 
         async def set_connected():
             """Set connected state."""
-            await xknx.connection_manager.connection_state_changed(
+            xknx.connection_manager.connection_state_changed(
                 XknxConnectionState.CONNECTED
             )
             assert self.main_thread != threading.get_ident()
@@ -111,6 +108,7 @@ class TestConnectionManager:
             # wait for side_effect to finish
             async with asyncio_timeout(1):
                 await xknx.connection_manager.connected.wait()
+                await asyncio.sleep(0)
             await xknx.stop()
 
     async def test_connection_information(self):
@@ -127,7 +125,7 @@ class TestConnectionManager:
         xknx.connection_manager.cemi_count_outgoing_error = 5
 
         # reset counters on new connection
-        await xknx.connection_manager.connection_state_changed(
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.CONNECTED, XknxConnectionType.TUNNEL_TCP
         )
         assert xknx.connection_manager.cemi_count_incoming == 0
@@ -142,7 +140,7 @@ class TestConnectionManager:
         xknx.connection_manager.cemi_count_outgoing = 5
         xknx.connection_manager.cemi_count_outgoing_error = 5
         # keep values until new connection; set connection timestamp to None
-        await xknx.connection_manager.connection_state_changed(
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
         assert xknx.connection_manager.cemi_count_incoming == 5
diff --git a/test/core_tests/state_updater_test.py b/test/core_tests/state_updater_test.py
index 1d4ab978..247e7a7f 100644
--- a/test/core_tests/state_updater_test.py
+++ b/test/core_tests/state_updater_test.py
@@ -200,15 +200,13 @@ class TestStateUpdater:
 
         assert xknx.state_updater.started
 
-        await xknx.connection_manager.connection_state_changed(
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
 
         assert not xknx.state_updater.started
 
-        await xknx.connection_manager.connection_state_changed(
-            XknxConnectionState.CONNECTED
-        )
+        xknx.connection_manager.connection_state_changed(XknxConnectionState.CONNECTED)
 
         assert xknx.state_updater.started
 
diff --git a/test/core_tests/task_registry_test.py b/test/core_tests/task_registry_test.py
index a808eec3..2bdfee4e 100644
--- a/test/core_tests/task_registry_test.py
+++ b/test/core_tests/task_registry_test.py
@@ -83,9 +83,7 @@ class TestTaskRegistry:
         xknx = XKNX()
         xknx.task_registry.start()
         assert len(xknx.connection_manager._connection_state_changed_cbs) == 1
-        await xknx.connection_manager.connection_state_changed(
-            XknxConnectionState.CONNECTED
-        )
+        xknx.connection_manager.connection_state_changed(XknxConnectionState.CONNECTED)
         # pylint: disable=attribute-defined-outside-init
         self.test = 0
 
@@ -106,16 +104,18 @@ class TestTaskRegistry:
         assert task._task is not None
         await time_travel(100)
         assert self.test == 1
-        await xknx.connection_manager.connection_state_changed(
+
+        xknx.connection_manager.connection_state_changed(
             XknxConnectionState.DISCONNECTED
         )
+        await asyncio.sleep(0)  # iterate loop to cancel task
         assert task._task is None
         assert self.test == 0
-        await xknx.connection_manager.connection_state_changed(
-            XknxConnectionState.CONNECTED
-        )
+
+        xknx.connection_manager.connection_state_changed(XknxConnectionState.CONNECTED)
         assert task._task is not None
         assert self.test == 0
+
         await time_travel(100)
         assert self.test == 1
         assert len(xknx.task_registry.tasks) == 1
diff --git a/test/xknx_test.py b/test/xknx_test.py
index 72e71f76..23076f50 100644
--- a/test/xknx_test.py
+++ b/test/xknx_test.py
@@ -1,7 +1,7 @@
 """Unit test for XKNX Module."""
 
 from pathlib import Path
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
@@ -37,7 +37,7 @@ class TestXknxModule:
 
     def test_register_connection_state_change_cb(self):
         """Test register con state callback."""
-        xknx = XKNX(connection_state_changed_cb=AsyncMock())
+        xknx = XKNX(connection_state_changed_cb=Mock())
         assert len(xknx.connection_manager._connection_state_changed_cbs) == 1
 
     @patch("xknx.io.KNXIPInterface._start", new_callable=AsyncMock)