Handle missing attrs in whois results (#65254)

* Handle missing attrs in whois results

- Some attrs are not set depending on where the
  domain is registered

- Fixes #65164

* Set to unknown instead of do not create

* no multi-line lambda
This commit is contained in:
J. Nick Koston 2022-01-30 15:19:04 -06:00 committed by GitHub
parent 473abb1793
commit 62fd31a1e7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 82 additions and 8 deletions

View file

@ -80,6 +80,13 @@ def _ensure_timezone(timestamp: datetime | None) -> datetime | None:
return timestamp return timestamp
def _fetch_attr_if_exists(domain: Domain, attr: str) -> str | None:
"""Fetch an attribute if it exists and is truthy or return None."""
if hasattr(domain, attr) and (value := getattr(domain, attr)):
return cast(str, value)
return None
SENSORS: tuple[WhoisSensorEntityDescription, ...] = ( SENSORS: tuple[WhoisSensorEntityDescription, ...] = (
WhoisSensorEntityDescription( WhoisSensorEntityDescription(
key="admin", key="admin",
@ -87,7 +94,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = (
icon="mdi:account-star", icon="mdi:account-star",
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=lambda domain: domain.admin if domain.admin else None, value_fn=lambda domain: _fetch_attr_if_exists(domain, "admin"),
), ),
WhoisSensorEntityDescription( WhoisSensorEntityDescription(
key="creation_date", key="creation_date",
@ -123,7 +130,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = (
icon="mdi:account", icon="mdi:account",
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=lambda domain: domain.owner if domain.owner else None, value_fn=lambda domain: _fetch_attr_if_exists(domain, "owner"),
), ),
WhoisSensorEntityDescription( WhoisSensorEntityDescription(
key="registrant", key="registrant",
@ -131,7 +138,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = (
icon="mdi:account-edit", icon="mdi:account-edit",
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=lambda domain: domain.registrant if domain.registrant else None, value_fn=lambda domain: _fetch_attr_if_exists(domain, "registrant"),
), ),
WhoisSensorEntityDescription( WhoisSensorEntityDescription(
key="registrar", key="registrar",
@ -147,7 +154,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = (
icon="mdi:store", icon="mdi:store",
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=lambda domain: domain.reseller if domain.reseller else None, value_fn=lambda domain: _fetch_attr_if_exists(domain, "reseller"),
), ),
) )
@ -190,7 +197,6 @@ async def async_setup_entry(
) )
for description in SENSORS for description in SENSORS
], ],
update_before_add=True,
) )

View file

@ -3,7 +3,7 @@ from __future__ import annotations
from collections.abc import Generator from collections.abc import Generator
from datetime import datetime from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest import pytest
@ -71,6 +71,33 @@ def mock_whois() -> Generator[MagicMock, None, None]:
yield whois_mock yield whois_mock
@pytest.fixture
def mock_whois_missing_some_attrs() -> Generator[Mock, None, None]:
"""Return a mocked query that only sets admin."""
class LimitedWhoisMock:
"""A limited mock of whois_query."""
def __init__(self, *args, **kwargs):
"""Mock only attributes the library always sets being available."""
self.creation_date = datetime(2019, 1, 1, 0, 0, 0)
self.dnssec = True
self.expiration_date = datetime(2023, 1, 1, 0, 0, 0)
self.last_updated = datetime(
2022, 1, 1, 0, 0, 0, tzinfo=dt_util.get_time_zone("Europe/Amsterdam")
)
self.name = "home-assistant.io"
self.name_servers = ["ns1.example.com", "ns2.example.com"]
self.registrar = "My Registrar"
self.status = "OK"
self.statuses = ["OK"]
with patch(
"homeassistant.components.whois.whois_query", LimitedWhoisMock
) as whois_mock:
yield whois_mock
@pytest.fixture @pytest.fixture
async def init_integration( async def init_integration(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_whois: MagicMock hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_whois: MagicMock
@ -84,6 +111,21 @@ async def init_integration(
return mock_config_entry return mock_config_entry
@pytest.fixture
async def init_integration_missing_some_attrs(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_whois_missing_some_attrs: MagicMock,
) -> MockConfigEntry:
"""Set up thewhois integration for testing."""
mock_config_entry.add_to_hass(hass)
await hass.config_entries.async_setup(mock_config_entry.entry_id)
await hass.async_block_till_done()
return mock_config_entry
@pytest.fixture @pytest.fixture
def enable_all_entities() -> Generator[AsyncMock, None, None]: def enable_all_entities() -> Generator[AsyncMock, None, None]:
"""Test fixture that ensures all entities are enabled in the registry.""" """Test fixture that ensures all entities are enabled in the registry."""

View file

@ -30,7 +30,7 @@ async def test_load_unload_config_entry(
await hass.async_block_till_done() await hass.async_block_till_done()
assert mock_config_entry.state is ConfigEntryState.LOADED assert mock_config_entry.state is ConfigEntryState.LOADED
assert len(mock_whois.mock_calls) == 2 assert len(mock_whois.mock_calls) == 1
await hass.config_entries.async_unload(mock_config_entry.entry_id) await hass.config_entries.async_unload(mock_config_entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -76,5 +76,5 @@ async def test_import_config(
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.config_entries.async_entries(DOMAIN)) == 1 assert len(hass.config_entries.async_entries(DOMAIN)) == 1
assert len(mock_whois.mock_calls) == 2 assert len(mock_whois.mock_calls) == 1
assert "the Whois platform in YAML is deprecated" in caplog.text assert "the Whois platform in YAML is deprecated" in caplog.text

View file

@ -143,6 +143,32 @@ async def test_whois_sensors(
assert device_entry.sw_version is None assert device_entry.sw_version is None
@pytest.mark.freeze_time("2022-01-01 12:00:00", tz_offset=0)
async def test_whois_sensors_missing_some_attrs(
hass: HomeAssistant,
enable_all_entities: AsyncMock,
init_integration_missing_some_attrs: MockConfigEntry,
) -> None:
"""Test the Whois sensors with owner and reseller missing."""
entity_registry = er.async_get(hass)
state = hass.states.get("sensor.home_assistant_io_last_updated")
entry = entity_registry.async_get("sensor.home_assistant_io_last_updated")
assert entry
assert state
assert entry.unique_id == "home-assistant.io_last_updated"
assert entry.entity_category == EntityCategory.DIAGNOSTIC
assert state.state == "2021-12-31T23:00:00+00:00"
assert state.attributes.get(ATTR_FRIENDLY_NAME) == "home-assistant.io Last Updated"
assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.TIMESTAMP
assert ATTR_ICON not in state.attributes
assert hass.states.get("sensor.home_assistant_io_owner").state == STATE_UNKNOWN
assert hass.states.get("sensor.home_assistant_io_reseller").state == STATE_UNKNOWN
assert hass.states.get("sensor.home_assistant_io_registrant").state == STATE_UNKNOWN
assert hass.states.get("sensor.home_assistant_io_admin").state == STATE_UNKNOWN
@pytest.mark.parametrize( @pytest.mark.parametrize(
"entity_id", "entity_id",
( (