Skip to content

Core Architecture: DataUpdateCoordinator & Diagnostics Integration #98

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Copilot wants to merge 7 commits into from
Closed

Core Architecture: DataUpdateCoordinator & Diagnostics Integration #98

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1ec476d
Initial plan
Copilot Aug 20, 2025
8c1cb6d
Enhanced diagnostics implementation with comprehensive data collectio…
Copilot Aug 20, 2025
ac73b41
Complete DataUpdateCoordinator improvements and test fixes
Copilot Aug 20, 2025
330a106
Update tests/test_coordinator_diagnostics.py
sirkirby Aug 25, 2025
6d7693e
Update custom_components/unifi_network_rules/utils/diagnostics.py
sirkirby Aug 25, 2025
9c39b9c
Update custom_components/unifi_network_rules/coordinator.py
sirkirby Aug 25, 2025
9529485
Fix missing regex constants and coordinator concurrency control issues
Copilot Aug 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 92 additions & 123 deletions custom_components/unifi_network_rules/coordinator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,7 @@ def __init__(
# Initialize config_entry to None - it will be looked up when needed
self.config_entry = None

# Update lock - prevent simultaneous updates
self._update_lock = asyncio.Lock()
# DataUpdateCoordinator handles concurrency internally, reducing need for manual locks

# Authentication state
self._authentication_in_progress = False
Expand Down Expand Up @@ -317,45 +316,40 @@ def _check_for_device_state_changes(self, previous_data: Dict[str, List[Any]], n
)

async def _async_update_data(self) -> Dict[str, List[Any]]:
"""Fetch data from API endpoint."""
# Use a lock to prevent concurrent updates, especially during authentication
if self._update_lock.locked():
LOGGER.debug("Another update is already in progress, waiting for it to complete")
# If an update is already in progress, wait for it to complete and use its result
if self.data:
return self.data
elif self._last_successful_data:
return self._last_successful_data
"""Fetch data from API endpoint.

async with self._update_lock:
try:
# Track authentication state at start of update
authentication_active = self._authentication_in_progress
if authentication_active:
LOGGER.warning("Update started while authentication is in progress - using cached data")
if self.data:
return self.data
elif self._last_successful_data:
return self._last_successful_data

# Proactively refresh the session to prevent 403 errors
# Only refresh every 5 minutes to avoid excessive API calls
refresh_interval = 300 # seconds
current_time = asyncio.get_event_loop().time()
last_refresh = getattr(self, "_last_session_refresh", 0)

if current_time - last_refresh > refresh_interval:
LOGGER.debug("Proactively refreshing session")
try:
# We'll track successful refreshes but not fail the update if refresh fails
refresh_success = await self.api.refresh_session()
if refresh_success:
self._last_session_refresh = current_time
LOGGER.debug("Session refresh successful")
else:
LOGGER.warning("Session refresh skipped or failed, continuing with update")
except Exception as refresh_err:
LOGGER.warning("Failed to refresh session: %s", str(refresh_err))
DataUpdateCoordinator provides built-in concurrency control for us,
eliminating the need for manual locks. We focus on UniFi-specific
authentication and error handling logic.
"""
try:
# Track authentication state at start of update
authentication_active = self._authentication_in_progress
if authentication_active:
LOGGER.warning("Update started while authentication is in progress - using cached data")
if self.data:
return self.data
elif self._last_successful_data:
return self._last_successful_data

# Proactively refresh the session to prevent 403 errors
# Only refresh every 5 minutes to avoid excessive API calls
refresh_interval = 300 # seconds
current_time = asyncio.get_event_loop().time()
last_refresh = getattr(self, "_last_session_refresh", 0)

if current_time - last_refresh > refresh_interval:
LOGGER.debug("Proactively refreshing session")
try:
# We'll track successful refreshes but not fail the update if refresh fails
refresh_success = await self.api.refresh_session()
if refresh_success:
self._last_session_refresh = current_time
LOGGER.debug("Session refresh successful")
else:
LOGGER.warning("Session refresh skipped or failed, continuing with update")
except Exception as refresh_err:
LOGGER.warning("Failed to refresh session: %s", str(refresh_err))

# Initialize with empty lists for each rule type
rules_data: Dict[str, List[Any]] = {
Expand Down Expand Up @@ -595,45 +589,45 @@ async def _async_update_data(self) -> Dict[str, List[Any]]:
len(self.vpn_servers),
len(self.devices))

return rules_data
return rules_data

except Exception as err:
LOGGER.error("Error updating coordinator data: %s", err)

# Check if this is an authentication error
auth_error = False
error_str = str(err).lower()
if "401 unauthorized" in error_str or "403 forbidden" in error_str:
auth_error = True
self._auth_failures += 1
self._authentication_in_progress = True
try:
LOGGER.warning("Authentication failure #%d during data update", self._auth_failures)

# Signal auth failure to entities
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_failure")

# Try to refresh the session if we haven't exceeded max failures
if self._auth_failures < self._max_auth_failures:
LOGGER.info("Attempting to refresh authentication session")
try:
await self.api.refresh_session(force=True)
# If we succeeded in refreshing, notify components
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_restored")
# Return the previous data
if self.data:
return self.data
except Exception as refresh_err:
LOGGER.error("Failed to refresh session: %s", refresh_err)
finally:
self._authentication_in_progress = False
except Exception as err:
LOGGER.error("Error updating coordinator data: %s", err)

# Return previous data during errors if available to prevent entity flickering
if self.data:
LOGGER.info("Returning previous data during error")
return self.data
# Check if this is an authentication error
auth_error = False
error_str = str(err).lower()
if "401 unauthorized" in error_str or "403 forbidden" in error_str:
auth_error = True
self._auth_failures += 1
self._authentication_in_progress = True
try:
LOGGER.warning("Authentication failure #%d during data update", self._auth_failures)

# Signal auth failure to entities
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_failure")

raise UpdateFailed(f"Error updating data: {err}")
# Try to refresh the session if we haven't exceeded max failures
if self._auth_failures < self._max_auth_failures:
LOGGER.info("Attempting to refresh authentication session")
try:
await self.api.refresh_session(force=True)
# If we succeeded in refreshing, notify components
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_restored")
# Return the previous data
if self.data:
return self.data
except Exception as refresh_err:
LOGGER.error("Failed to refresh session: %s", refresh_err)
finally:
self._authentication_in_progress = False

# Return previous data during errors if available to prevent entity flickering
if self.data:
LOGGER.info("Returning previous data during error")
return self.data

raise UpdateFailed(f"Error updating data: {err}")

def _check_for_deleted_rules(self, new_data: Dict[str, List[Any]]) -> None:
"""Check for rules previously known but not in the new data, and trigger their removal."""
Expand Down Expand Up @@ -964,68 +958,40 @@ def _handle_websocket_message(self, message: dict[str, Any]) -> None:
pass

async def _controlled_refresh_wrapper(self):
"""Wrapper for the controlled refresh process to ensure proper semaphore handling."""
# Acquire semaphore before starting refresh
async def controlled_refresh():
async with self._refresh_semaphore:
await self._force_refresh_with_cache_clear()
# Wait a moment before refreshing entities to let states settle
await asyncio.sleep(0.5)
self.async_update_listeners()
"""Request a refresh using DataUpdateCoordinator's built-in mechanisms.

# Start the controlled refresh task and await it
await controlled_refresh()
Uses coordinator's throttling and failure handling instead of custom semaphore.
"""
# Use DataUpdateCoordinator's built-in refresh request
# This provides better resilience and throttling
await self.async_request_refresh()

async def _force_refresh_with_cache_clear(self) -> None:
"""Force a refresh with cache clearing to ensure fresh data.

This method is triggered by WebSocket events and follows the same core refresh
and entity management logic as the regular polling updates:
1. Clear API cache (but preserve authentication)
2. Call async_refresh() which updates rule collections
3. Process new entities through the same entity creation path
4. Check for deleted rules to maintain consistency with polling
"""Force a refresh with cache clearing using DataUpdateCoordinator mechanisms.

The only major difference is the explicit call to _check_for_deleted_rules()
which happens automatically during polling updates.
Leverages the coordinator's built-in refresh with custom cache clearing for UniFi.
"""
try:
# Log that we're starting a refresh
log_websocket("Starting forced refresh after rule change detected")

# Store previous data for deletion detection
previous_data = self.data.copy() if self.data else {}

# Clear the API cache to ensure we get fresh data
# But do so without disrupting authentication
LOGGER.debug("Clearing API cache")
if hasattr(self.api, "clear_cache"):
await self.api.clear_cache() # Modified in API to preserve auth
await self.api.clear_cache()
else:
LOGGER.warning("API object does not have clear_cache method")
LOGGER.debug("API cache cleared")
log_data("Cache cleared before refresh")

# Force a full data refresh
refresh_successful = await self.async_refresh()
# Use DataUpdateCoordinator's built-in refresh mechanism
await self.async_request_refresh()

log_data("Refresh completed successfully after WebSocket event")

if refresh_successful:
# After refreshing data, discovery and deletion checks are handled within async_refresh -> _async_update_data
# REMOVED: await self.process_new_entities() # Redundant
# REMOVED: if previous_data: # Incorrect check
# REMOVED: self._check_for_deleted_rules(previous_data) # Incorrect check

# Update the data timestamp
self._last_update = self.hass.loop.time()

# Force an update of all entities
self.async_update_listeners()

log_data("Refresh completed successfully after WebSocket event")
else:
LOGGER.error("WebSocket-triggered refresh failed")
except Exception as err:
LOGGER.error("Error during forced refresh: %s", err)
# Let coordinator handle the error with its built-in mechanisms

@callback
def shutdown(self) -> None:
Expand All @@ -1042,7 +1008,10 @@ async def async_shutdown(self) -> None:
# For example, wait for pending tasks to complete

async def _handle_auth_failure(self):
"""Handle authentication failures from API operations."""
"""Handle authentication failures using DataUpdateCoordinator patterns.

Leverages coordinator's built-in failure handling and backoff mechanisms.
"""
LOGGER.info("Authentication failure callback triggered, requesting data refresh")
# Reset authentication flag
self._authentication_in_progress = False
Expand All @@ -1053,8 +1022,8 @@ async def _handle_auth_failure(self):
# Notify any entities that have optimistic state to handle auth issues appropriately
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_failure")

# Request a refresh with some delay to allow auth to stabilize
await asyncio.sleep(2.0)
# Brief delay to allow auth to stabilize (coordinator handles longer backoff)
await asyncio.sleep(1.0)

# Ensure a fresh session before refresh
try:
Expand All @@ -1064,8 +1033,8 @@ async def _handle_auth_failure(self):
except Exception as err:
LOGGER.error("Error refreshing session during auth recovery: %s", str(err))

# Force a full refresh
await self.async_refresh()
# Use DataUpdateCoordinator's refresh mechanism instead of direct async_refresh
await self.async_request_refresh()

# After successful refresh, notify components to clear any error states
async_dispatcher_send(self.hass, f"{DOMAIN}_auth_restored")
Expand Down
16 changes: 16 additions & 0 deletions custom_components/unifi_network_rules/diagnostics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
"""Diagnostics support for UniFi Network Rules."""
from __future__ import annotations

from typing import Any, Dict

from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant

from .utils.diagnostics import async_get_config_entry_diagnostics as _async_get_config_entry_diagnostics


async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry
) -> Dict[str, Any]:
"""Return diagnostics for a config entry."""
return await _async_get_config_entry_diagnostics(hass, entry)
Loading