Skip to content

Commit

Permalink
refactor(switch,core): Improve state verification and init data fetching
Browse files Browse the repository at this point in the history 
- refactor(switch): Improve state verification and error handling for network switches
- Enhanced state verification mechanism for traffic routes and firewall policies
- Directly fetch fresh data from API during state verification
- Added more robust error handling and logging during state changes
- Removed redundant state tracking attributes
- Simplified state verification and toggle logic
- Improved logging for switch state transitions and verification attempts
- refactor(core): Improve data fetching and error handling in UniFi Network integration
- Enhanced data update coordinator with more robust error handling
- Added explicit UpdateFailed exception for data retrieval errors
- Implemented initial data fetch with comprehensive logging
- Improved error checking and data validation during setup
- Added debug logging for firewall policies and traffic routes count
- Ensured fallback to empty lists when no data is retrieved

Everything works - but no firewall rules or other entities shows up #18
  • Loading branch information
@sirkirby
sirkirby authored Feb 6, 2025
1 parent 1805723 commit 05042e6
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 69 deletions.
36 changes: 29 additions & 7 deletions custom_components/unifi_network_rules/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from homeassistant.core import HomeAssistant
from homeassistant.const import CONF_HOST, CONF_USERNAME, CONF_PASSWORD
from homeassistant.helpers.typing import ConfigType
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from homeassistant.helpers import config_validation as cv
from homeassistant.exceptions import ConfigEntryNotReady

Expand Down Expand Up @@ -63,23 +63,29 @@ async def async_update_data():
"""Fetch data from API."""
_LOGGER.debug("Starting data update")
try:
# Get firewall policies
policies_success, policies, policies_error = await api.get_firewall_policies()
if not policies_success:
_LOGGER.error(f"Failed to fetch policies: {policies_error}")
raise Exception(f"Failed to fetch policies: {policies_error}")
raise UpdateFailed(f"Failed to fetch policies: {policies_error}")

_LOGGER.debug(f"Retrieved {len(policies) if policies else 0} firewall policies")

# Get traffic routes
routes_success, routes, routes_error = await api.get_traffic_routes()
if not routes_success:
_LOGGER.error(f"Failed to fetch routes: {routes_error}")
raise Exception(f"Failed to fetch routes: {routes_error}")
raise UpdateFailed(f"Failed to fetch routes: {routes_error}")

_LOGGER.debug(f"Retrieved {len(routes) if routes else 0} traffic routes")

return {
"firewall_policies": policies,
"traffic_routes": routes
"firewall_policies": policies or [],
"traffic_routes": routes or []
}
except Exception as e:
_LOGGER.exception("Error in update_data")
raise
raise UpdateFailed(f"Data update failed: {str(e)}")

_LOGGER.debug("Creating coordinator")
coordinator = DataUpdateCoordinator(
Expand All @@ -90,14 +96,30 @@ async def async_update_data():
update_interval=timedelta(minutes=update_interval),
)

# Perform initial data fetch
_LOGGER.debug("Performing initial data fetch")
await coordinator.async_config_entry_first_refresh()

# Verify we have data before proceeding
if not coordinator.data:
error_msg = "No data received from UniFi Network during setup"
_LOGGER.error(error_msg)
await api.cleanup()
raise ConfigEntryNotReady(error_msg)

# Log the initial data for debugging
_LOGGER.debug("Initial coordinator data:")
_LOGGER.debug("Firewall Policies: %d", len(coordinator.data.get("firewall_policies", [])))
_LOGGER.debug("Traffic Routes: %d", len(coordinator.data.get("traffic_routes", [])))

# Store API and coordinator
_LOGGER.debug("Storing API and coordinator")
hass.data[DOMAIN][entry.entry_id] = {
'api': api,
'coordinator': coordinator,
}

# Set up platform first
# Set up platform
_LOGGER.debug("Setting up platform")
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

Expand Down
148 changes: 86 additions & 62 deletions custom_components/unifi_network_rules/switch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
entities.append(UDMTrafficRouteSwitch(coordinator, api, route))

async_add_entities(entities)

class UDMTrafficRouteSwitch(CoordinatorEntity, SwitchEntity):
"""Representation of a UDM Traffic Route Switch."""

Expand All @@ -45,7 +45,6 @@ def __init__(self, coordinator, api, route):
self._route_id = route['_id']
self._route = route
self._pending_state = None
self._state_update_time = None
_LOGGER.info("Initialized traffic route switch: %s (ID: %s)", self._attr_name, self._route_id)

@callback
Expand Down Expand Up @@ -81,9 +80,6 @@ def is_on(self) -> bool:
"""Return true if the switch is on."""
if self._pending_state is not None:
return self._pending_state
if not self._route:
_LOGGER.warning("No route data available for %s", self._attr_name)
return False
return bool(self._route.get('enabled', False))

async def async_turn_on(self, **kwargs: Any) -> None:
Expand All @@ -99,24 +95,49 @@ async def async_turn_off(self, **kwargs: Any) -> None:
async def _verify_state_change(self, target_state: bool, max_attempts: int = 3) -> bool:
"""Verify that the state change was successful."""
for attempt in range(max_attempts):
await self.coordinator.async_request_refresh()
if self._route.get('enabled') == target_state:
_LOGGER.info("State change verified for %s", self._attr_name)
return True
_LOGGER.warning(
"State verification attempt %d failed for %s. Expected: %s, Got: %s",
attempt + 1,
self._attr_name,
target_state,
self._route.get('enabled')
)
await asyncio.sleep(2)
try:
# Get fresh data directly from the API instead of coordinator
success, routes, error = await self._api.get_traffic_routes()
if not success:
_LOGGER.error("Failed to fetch routes for verification: %s", error)
await asyncio.sleep(2)
continue

current_route = next((r for r in routes if r['_id'] == self._route_id), None)
if not current_route:
_LOGGER.error("Route not found during verification")
await asyncio.sleep(2)
continue

current_state = current_route.get('enabled', False)
if current_state == target_state:
# Update the coordinator data to reflect the new state
await self.coordinator.async_request_refresh()
_LOGGER.info("State change verified for %s", self._attr_name)
return True

_LOGGER.warning(
"State verification attempt %d/%d failed for %s. Expected: %s, Got: %s",
attempt + 1,
max_attempts,
self._attr_name,
target_state,
current_state
)
await asyncio.sleep(2)
except Exception as e:
_LOGGER.error("Error during state verification: %s", str(e))
await asyncio.sleep(2)

return False

async def _toggle(self, new_state: bool) -> None:
"""Toggle the route state."""
try:
_LOGGER.info("Current route state: %s", self._route)
if self._route.get('enabled') == new_state:
_LOGGER.info("%s is already in desired state: %s", self._attr_name, new_state)
return

self._pending_state = new_state
self.async_write_ha_state()

Expand All @@ -125,7 +146,9 @@ async def _toggle(self, new_state: bool) -> None:
if success:
_LOGGER.info("API reports successful toggle for %s", self._attr_name)

# Verify the state change
# Add a small delay before verification to allow for API propagation
await asyncio.sleep(1)

if await self._verify_state_change(new_state):
_LOGGER.info("Successfully verified state change for %s", self._attr_name)
else:
Expand Down Expand Up @@ -157,24 +180,14 @@ def __init__(self, coordinator, api, policy):
self._policy_id = policy['_id']
self._policy = policy
self._pending_state = None
self._state_update_time = None
_LOGGER.info("Initialized firewall policy switch: %s (ID: %s)", self._attr_name, self._policy_id)

def get_policy(self):
"""Get the current policy from coordinator data."""
if not self.coordinator.data:
_LOGGER.debug("No coordinator data available for %s", self._attr_name)
return None
policies = self.coordinator.data.get('firewall_policies', [])
policy = next((p for p in policies if p['_id'] == self._policy_id), None)
if policy is None:
_LOGGER.debug("Policy not found in coordinator data for %s", self._attr_name)
return policy

@callback
def _handle_coordinator_update(self) -> None:
"""Handle updated data from the coordinator."""
new_policy = self.get_policy()
policies = self.coordinator.data.get('firewall_policies', [])
new_policy = next((p for p in policies if p['_id'] == self._policy_id), None)

if new_policy:
if self._policy.get('enabled') != new_policy.get('enabled'):
_LOGGER.info(
Expand Down Expand Up @@ -202,9 +215,6 @@ def is_on(self) -> bool:
"""Return true if the switch is on."""
if self._pending_state is not None:
return self._pending_state
if not self._policy:
_LOGGER.warning("No policy data available for %s", self._attr_name)
return False
return bool(self._policy.get('enabled', False))

async def async_turn_on(self, **kwargs: Any) -> None:
Expand All @@ -220,38 +230,49 @@ async def async_turn_off(self, **kwargs: Any) -> None:
async def _verify_state_change(self, target_state: bool, max_attempts: int = 3) -> bool:
"""Verify that the state change was successful."""
for attempt in range(max_attempts):
await self.coordinator.async_request_refresh()
current_state = self._policy.get('enabled', False)

if current_state == target_state:
_LOGGER.info("State change verified for %s - Target: %s, Current: %s",
self._attr_name, target_state, current_state)
return True

_LOGGER.warning(
"State verification attempt %d/%d failed for %s. Expected: %s, Got: %s",
attempt + 1,
max_attempts,
self._attr_name,
target_state,
current_state
)
await asyncio.sleep(2)

# Only return False if we've exhausted all attempts and the states don't match
try:
# Get fresh data directly from the API instead of coordinator
success, policies, error = await self._api.get_firewall_policies()
if not success:
_LOGGER.error("Failed to fetch policies for verification: %s", error)
await asyncio.sleep(2)
continue

current_policy = next((p for p in policies if p['_id'] == self._policy_id), None)
if not current_policy:
_LOGGER.error("Policy not found during verification")
await asyncio.sleep(2)
continue

current_state = current_policy.get('enabled', False)
if current_state == target_state:
# Update the coordinator data to reflect the new state
await self.coordinator.async_request_refresh()
_LOGGER.info("State change verified for %s", self._attr_name)
return True

_LOGGER.warning(
"State verification attempt %d/%d failed for %s. Expected: %s, Got: %s",
attempt + 1,
max_attempts,
self._attr_name,
target_state,
current_state
)
await asyncio.sleep(2)
except Exception as e:
_LOGGER.error("Error during state verification: %s", str(e))
await asyncio.sleep(2)

return False

async def _toggle(self, new_state: bool) -> None:
"""Toggle the policy state."""
try:
_LOGGER.info("Attempting to toggle %s to %s", self._attr_name, new_state)
current_state = self._policy.get('enabled', False)

# If the current state already matches the target state, no need to toggle
if current_state == new_state:
if self._policy.get('enabled') == new_state:
_LOGGER.info("%s is already in desired state: %s", self._attr_name, new_state)
return

self._pending_state = new_state
self.async_write_ha_state()

Expand All @@ -260,13 +281,16 @@ async def _toggle(self, new_state: bool) -> None:
if success:
_LOGGER.info("API reports successful toggle for %s", self._attr_name)

# Verify the state change
# Add a small delay before verification to allow for API propagation
await asyncio.sleep(1)

if await self._verify_state_change(new_state):
_LOGGER.info("Successfully verified state change for %s", self._attr_name)
else:
self._pending_state = None
raise HomeAssistantError(
f"Failed to verify state change for {self._attr_name} after multiple attempts"
f"Failed to verify state change for {self._attr_name}. "
f"Target state: {new_state}, Current state: {self._policy.get('enabled')}"
)
else:
self._pending_state = None
Expand Down

0 comments on commit 05042e6

Please sign in to comment.