Ensure bleak_retry_connector uses HaBleakClientWrapper (#79132)

This commit is contained in:
J. Nick Koston 2022-09-27 08:06:10 -10:00 committed by GitHub
parent 8eaa22cf90
commit 5eb50f63fd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 49 additions and 0 deletions

View file

@ -3,20 +3,39 @@
from __future__ import annotations
import bleak
from bleak.backends.service import BleakGATTServiceCollection
import bleak_retry_connector
from .models import HaBleakClientWrapper, HaBleakScannerWrapper
ORIGINAL_BLEAK_SCANNER = bleak.BleakScanner
ORIGINAL_BLEAK_CLIENT = bleak.BleakClient
ORIGINAL_BLEAK_RETRY_CONNECTOR_CLIENT = (
bleak_retry_connector.BleakClientWithServiceCache
)
def install_multiple_bleak_catcher() -> None:
"""Wrap the bleak classes to return the shared instance if multiple instances are detected."""
bleak.BleakScanner = HaBleakScannerWrapper # type: ignore[misc, assignment]
bleak.BleakClient = HaBleakClientWrapper # type: ignore[misc]
bleak_retry_connector.BleakClientWithServiceCache = HaBleakClientWithServiceCache # type: ignore[misc,assignment]
def uninstall_multiple_bleak_catcher() -> None:
"""Unwrap the bleak classes."""
bleak.BleakScanner = ORIGINAL_BLEAK_SCANNER # type: ignore[misc]
bleak.BleakClient = ORIGINAL_BLEAK_CLIENT # type: ignore[misc]
bleak_retry_connector.BleakClientWithServiceCache = ORIGINAL_BLEAK_RETRY_CONNECTOR_CLIENT # type: ignore[misc]
class HaBleakClientWithServiceCache(HaBleakClientWrapper):
"""A BleakClient that implements service caching."""
def set_cached_services(self, services: BleakGATTServiceCollection | None) -> None:
"""Set the cached services.
No longer used since bleak 0.17+ has service caching built-in.
This was only kept for backwards compatibility.
"""