From 7af1fb02c6dff60ea793c9a4fb5bd900286ad9ae Mon Sep 17 00:00:00 2001 From: Patrick762 Date: Mon, 31 Jul 2023 19:01:56 +0200 Subject: [PATCH] separated coordinator and sensor --- custom_components/bluetti_bt/__init__.py | 12 +- custom_components/bluetti_bt/coordinator.py | 153 ++++++++++++++++++++ custom_components/bluetti_bt/sensor.py | 149 +------------------ 3 files changed, 168 insertions(+), 146 deletions(-) create mode 100644 custom_components/bluetti_bt/coordinator.py diff --git a/custom_components/bluetti_bt/__init__.py b/custom_components/bluetti_bt/__init__.py index 99458b0..798a5cf 100644 --- a/custom_components/bluetti_bt/__init__.py +++ b/custom_components/bluetti_bt/__init__.py @@ -6,12 +6,13 @@ import re from homeassistant.components import bluetooth from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_ADDRESS, CONF_TYPE, Platform +from homeassistant.const import CONF_ADDRESS, CONF_TYPE, CONF_NAME, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers.entity import DeviceInfo from homeassistant.exceptions import ConfigEntryNotReady from .const import DOMAIN, MANUFACTURER +from .coordinator import PollingCoordinator PLATFORMS: [Platform] = [Platform.SENSOR] @@ -20,6 +21,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Bluetti Powerstation from a config entry.""" address = entry.data.get(CONF_ADDRESS) + device_name = entry.data.get(CONF_NAME) if address is None: return False @@ -27,6 +29,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if not bluetooth.async_address_present(hass, address): raise ConfigEntryNotReady("Bluetti device not present") + # Create coordinator for polling + coordinator = PollingCoordinator(hass, address, device_name) + await coordinator.async_config_entry_first_refresh() + + # Create data structure + hass.data.setdefault(DOMAIN, {}) + hass.data[DOMAIN].setdefault(entry.entry_id, coordinator) + # Setup platforms await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) diff --git a/custom_components/bluetti_bt/coordinator.py b/custom_components/bluetti_bt/coordinator.py new file mode 100644 index 0000000..e873948 --- /dev/null +++ b/custom_components/bluetti_bt/coordinator.py @@ -0,0 +1,153 @@ +"""Coordinator for Bluetti integration.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +import logging +from typing import cast +import async_timeout + +from bleak import BleakClient, BleakError + +from homeassistant.components import bluetooth +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import ( + DataUpdateCoordinator, +) + +from bluetti_mqtt.bluetooth.client import BluetoothClient +from bluetti_mqtt.bluetooth import ( + BadConnectionError, + ModbusError, + ParseError, + build_device, +) + +_LOGGER = logging.getLogger(__name__) + + +class PollingCoordinator(DataUpdateCoordinator): + """Polling coordinator.""" + + def __init__(self, hass: HomeAssistant, address, device_name: str): + """Initialize coordinator.""" + super().__init__( + hass, + _LOGGER, + name="Bluetti polling coordinator", + update_interval=timedelta(seconds=20), + ) + self._address = address + self.notify_future = None + self.current_command = None + self.notify_response = bytearray() + self.bluetti_device = build_device(address, device_name) + + async def _async_update_data(self): + """Fetch data from API endpoint. + + This is the place to pre-process the data to lookup tables + so entities can quickly look up their data. + """ + self.logger.debug("Polling data") + + device = bluetooth.async_ble_device_from_address(self.hass, self._address) + if device is None: + self.logger.error("Device not available") + return None + + if self.bluetti_device is None: + self.logger.error("Device type not found") + return None + + # Polling + client = BleakClient(device) + parsed_data: dict = {} + + try: + async with async_timeout.timeout(15): + await client.connect() + + await client.start_notify( + BluetoothClient.NOTIFY_UUID, self._notification_handler + ) + + for command in self.bluetti_device.polling_commands: + try: + # Prepare to make request + self.current_command = command + self.notify_future = self.hass.loop.create_future() + self.notify_response = bytearray() + + # Make request + self.logger.debug("Requesting %s", command) + await client.write_gatt_char( + BluetoothClient.WRITE_UUID, bytes(command) + ) + + # Wait for response + res = await asyncio.wait_for( + self.notify_future, timeout=BluetoothClient.RESPONSE_TIMEOUT + ) + + # Process data + self.logger.debug("Got %s bytes", len(res)) + response = cast(bytes, res) + body = command.parse_response(response) + parsed = self.bluetti_device.parse( + command.starting_address, body + ) + + self.logger.warning("Parsed data: %s", parsed) + parsed_data.update(parsed) + + except TimeoutError: + self.logger.error("Polling timed out (address: %s)", self._address) + except ParseError: + self.logger.debug("Got a parse exception...") + except ModbusError as err: + self.logger.debug( + "Got an invalid request error for %s: %s", + command, + err, + ) + except (BadConnectionError, BleakError) as err: + self.logger.debug("Needed to disconnect due to error: %s", err) + except TimeoutError: + self.logger.error("Polling timed out") + return None + except BleakError as err: + self.logger.error("Bleak error: %s", err) + return None + finally: + await client.disconnect() + + # Pass data back to sensors + return parsed_data + + def _notification_handler(self, _sender: int, data: bytearray): + """Handle bt data.""" + + # Ignore notifications we don't expect + if not self.notify_future or self.notify_future.done(): + return + + # If something went wrong, we might get weird data. + if data == b"AT+NAME?\r" or data == b"AT+ADV?\r": + err = BadConnectionError("Got AT+ notification") + self.notify_future.set_exception(err) + return + + # Save data + self.notify_response.extend(data) + + if len(self.notify_response) == self.current_command.response_size(): + if self.current_command.is_valid_response(self.notify_response): + self.notify_future.set_result(self.notify_response) + else: + self.notify_future.set_exception(ParseError("Failed checksum")) + elif self.current_command.is_exception_response(self.notify_response): + # We got a MODBUS command exception + msg = f"MODBUS Exception {self.current_command}: {self.notify_response[2]}" + self.notify_future.set_exception(ModbusError(msg)) diff --git a/custom_components/bluetti_bt/sensor.py b/custom_components/bluetti_bt/sensor.py index 1a6e746..088b8fc 100644 --- a/custom_components/bluetti_bt/sensor.py +++ b/custom_components/bluetti_bt/sensor.py @@ -2,13 +2,7 @@ from __future__ import annotations -import asyncio -from datetime import timedelta import logging -from typing import cast -import async_timeout - -from bleak import BleakClient, BleakError from homeassistant.components import bluetooth from homeassistant.components.sensor import SensorEntity, CONF_STATE_CLASS @@ -25,23 +19,18 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.update_coordinator import ( CoordinatorEntity, - DataUpdateCoordinator, ) -from bluetti_mqtt.bluetooth.client import BluetoothClient -from bluetti_mqtt.bluetooth import ( - BadConnectionError, - ModbusError, - ParseError, - build_device, -) +from bluetti_mqtt.bluetooth import build_device from bluetti_mqtt.mqtt_client import NORMAL_DEVICE_FIELDS from . import device_info as dev_info, get_unique_id from .const import ( + DOMAIN, API_RESPONSE_BATTERY_RANGE_START, API_RESPONSE_BATTERY_RANGE_END, ) +from .coordinator import PollingCoordinator _LOGGER = logging.getLogger(__name__) @@ -56,10 +45,6 @@ async def async_setup_entry( if address is None: _LOGGER.error("Device has no address") - # Create coordinator for polling - coordinator = PollingCoordinator(hass, address, device_name) - await coordinator.async_config_entry_first_refresh() - # Generate device info _LOGGER.info("Creating sensors for device with address %s", address) device_info = dev_info(entry) @@ -80,7 +65,7 @@ async def async_setup_entry( sensors_to_add.append( BluettiSensor( - coordinator, + hass.data[DOMAIN][entry.entry_id], device_info, address, field_key, @@ -95,132 +80,6 @@ async def async_setup_entry( async_add_entities(sensors_to_add) -class PollingCoordinator(DataUpdateCoordinator): - """Polling coordinator.""" - - def __init__(self, hass: HomeAssistant, address, device_name: str): - """Initialize coordinator.""" - super().__init__( - hass, - _LOGGER, - name="Bluetti polling coordinator", - update_interval=timedelta(seconds=20), - ) - self._address = address - self.notify_future = None - self.current_command = None - self.notify_response = bytearray() - self.bluetti_device = build_device(address, device_name) - - async def _async_update_data(self): - """Fetch data from API endpoint. - - This is the place to pre-process the data to lookup tables - so entities can quickly look up their data. - """ - self.logger.debug("Polling data") - - device = bluetooth.async_ble_device_from_address(self.hass, self._address) - if device is None: - self.logger.error("Device not available") - return None - - if self.bluetti_device is None: - self.logger.error("Device type not found") - return None - - # Polling - client = BleakClient(device) - parsed_data: dict = {} - - try: - async with async_timeout.timeout(15): - await client.connect() - - await client.start_notify( - BluetoothClient.NOTIFY_UUID, self._notification_handler - ) - - for command in self.bluetti_device.polling_commands: - try: - # Prepare to make request - self.current_command = command - self.notify_future = self.hass.loop.create_future() - self.notify_response = bytearray() - - # Make request - self.logger.debug("Requesting %s", command) - await client.write_gatt_char( - BluetoothClient.WRITE_UUID, bytes(command) - ) - - # Wait for response - res = await asyncio.wait_for( - self.notify_future, timeout=BluetoothClient.RESPONSE_TIMEOUT - ) - - # Process data - self.logger.debug("Got %s bytes", len(res)) - response = cast(bytes, res) - body = command.parse_response(response) - parsed = self.bluetti_device.parse( - command.starting_address, body - ) - - self.logger.warning("Parsed data: %s", parsed) - parsed_data.update(parsed) - - except TimeoutError: - self.logger.error("Polling timed out") - except ParseError: - self.logger.debug("Got a parse exception...") - except ModbusError as err: - self.logger.debug( - "Got an invalid request error for %s: %s", - command, - err, - ) - except (BadConnectionError, BleakError) as err: - self.logger.debug("Needed to disconnect due to error: %s", err) - except TimeoutError: - self.logger.error("Polling timed out") - return None - except BleakError as err: - self.logger.error("Bleak error: %s", err) - return None - finally: - await client.disconnect() - - # Pass data back to sensors - return parsed_data - - def _notification_handler(self, _sender: int, data: bytearray): - """Handle bt data.""" - - # Ignore notifications we don't expect - if not self.notify_future or self.notify_future.done(): - return - - # If something went wrong, we might get weird data. - if data == b"AT+NAME?\r" or data == b"AT+ADV?\r": - err = BadConnectionError("Got AT+ notification") - self.notify_future.set_exception(err) - return - - # Save data - self.notify_response.extend(data) - - if len(self.notify_response) == self.current_command.response_size(): - if self.current_command.is_valid_response(self.notify_response): - self.notify_future.set_result(self.notify_response) - else: - self.notify_future.set_exception(ParseError("Failed checksum")) - elif self.current_command.is_exception_response(self.notify_response): - # We got a MODBUS command exception - msg = f"MODBUS Exception {self.current_command}: {self.notify_response[2]}" - self.notify_future.set_exception(ModbusError(msg)) - - class BluettiSensor(CoordinatorEntity, SensorEntity): """Bluetti universal sensor."""