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:
parent
473abb1793
commit
62fd31a1e7
4 changed files with 82 additions and 8 deletions
|
@ -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,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
@ -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."""
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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",
|
||||||
(
|
(
|
||||||
|
|
Loading…
Add table
Reference in a new issue