diff --git a/homeassistant/components/whois/sensor.py b/homeassistant/components/whois/sensor.py index 2cbae147a78..0459651e693 100644 --- a/homeassistant/components/whois/sensor.py +++ b/homeassistant/components/whois/sensor.py @@ -80,6 +80,13 @@ def _ensure_timezone(timestamp: datetime | None) -> datetime | None: 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, ...] = ( WhoisSensorEntityDescription( key="admin", @@ -87,7 +94,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = ( icon="mdi:account-star", entity_category=EntityCategory.DIAGNOSTIC, 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( key="creation_date", @@ -123,7 +130,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = ( icon="mdi:account", entity_category=EntityCategory.DIAGNOSTIC, 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( key="registrant", @@ -131,7 +138,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = ( icon="mdi:account-edit", entity_category=EntityCategory.DIAGNOSTIC, 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( key="registrar", @@ -147,7 +154,7 @@ SENSORS: tuple[WhoisSensorEntityDescription, ...] = ( icon="mdi:store", entity_category=EntityCategory.DIAGNOSTIC, 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 ], - update_before_add=True, ) diff --git a/tests/components/whois/conftest.py b/tests/components/whois/conftest.py index bbda3b101f5..ef14750356e 100644 --- a/tests/components/whois/conftest.py +++ b/tests/components/whois/conftest.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Generator from datetime import datetime -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -71,6 +71,33 @@ def mock_whois() -> Generator[MagicMock, None, None]: 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 async def init_integration( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_whois: MagicMock @@ -84,6 +111,21 @@ async def init_integration( 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 def enable_all_entities() -> Generator[AsyncMock, None, None]: """Test fixture that ensures all entities are enabled in the registry.""" diff --git a/tests/components/whois/test_init.py b/tests/components/whois/test_init.py index 6701d8ed20c..3cd9efc801d 100644 --- a/tests/components/whois/test_init.py +++ b/tests/components/whois/test_init.py @@ -30,7 +30,7 @@ async def test_load_unload_config_entry( await hass.async_block_till_done() 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.async_block_till_done() @@ -76,5 +76,5 @@ async def test_import_config( await hass.async_block_till_done() 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 diff --git a/tests/components/whois/test_sensor.py b/tests/components/whois/test_sensor.py index b0e9862eb64..e824522ed09 100644 --- a/tests/components/whois/test_sensor.py +++ b/tests/components/whois/test_sensor.py @@ -143,6 +143,32 @@ async def test_whois_sensors( 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( "entity_id", (