Whois component maintenance, PR 1 of 3: Light refactoring, adding abstraction

homeassistant/components/whois/__init__.py

@@ -2,20 +2,14 @@
 
 from __future__ import annotations
 
-from whois import Domain, query as whois_query
-from whois.exceptions import (
-    FailedParsingWhoisOutput,
-    UnknownDateFormat,
-    UnknownTld,
-    WhoisCommandFailed,
-)
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
+from . import helper
 from .const import DOMAIN, LOGGER, PLATFORMS, SCAN_INTERVAL
+from .helper import Domain
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -25,11 +19,11 @@ async def _async_query_domain() -> Domain | None:
         """Query WHOIS for domain information."""
         try:
             return await hass.async_add_executor_job(
-                whois_query, entry.data[CONF_DOMAIN]
+                helper.query, entry.data[CONF_DOMAIN]
             )
-        except UnknownTld as ex:
+        except helper.WhoisUnknownTLD as ex:
             raise UpdateFailed("Could not set up whois, TLD is unknown") from ex
-        except (FailedParsingWhoisOutput, WhoisCommandFailed, UnknownDateFormat) as ex:
+        except helper.GenericWhoisError as ex:
             raise UpdateFailed("An error occurred during WHOIS lookup") from ex
 
     coordinator: DataUpdateCoordinator[Domain | None] = DataUpdateCoordinator(

homeassistant/components/whois/config_flow.py

@@ -5,17 +5,11 @@
 from typing import Any
 
 import voluptuous as vol
-import whois
-from whois.exceptions import (
-    FailedParsingWhoisOutput,
-    UnknownDateFormat,
-    UnknownTld,
-    WhoisCommandFailed,
-)
 
 from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
 from homeassistant.const import CONF_DOMAIN
 
+from . import helper
 from .const import DOMAIN
 
 
@@ -39,15 +33,11 @@ async def async_step_user(
             self._abort_if_unique_id_configured()
 
             try:
-                await self.hass.async_add_executor_job(whois.query, domain)
-            except UnknownTld:
+                await self.hass.async_add_executor_job(helper.query, domain)
+            except helper.WhoisUnknownTLD:
                 errors["base"] = "unknown_tld"
-            except WhoisCommandFailed:
+            except helper.GenericWhoisError:
                 errors["base"] = "whois_command_failed"
-            except FailedParsingWhoisOutput:
-                errors["base"] = "unexpected_response"
-            except UnknownDateFormat:
-                errors["base"] = "unknown_date_format"
             else:
                 return self.async_create_entry(
                     title=self.imported_name or user_input[CONF_DOMAIN],

homeassistant/components/whois/diagnostics.py

@@ -4,13 +4,12 @@
 
 from typing import Any
 
-from whois import Domain
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN
+from .helper import Domain
 
 
 async def async_get_config_entry_diagnostics(

homeassistant/components/whois/helper.py

@@ -0,0 +1,64 @@
+"""A helper class that abstracts away the usage of external whois libraries."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import whois
+
+
+@dataclass(kw_only=True)
+class Domain:
+    """A class internally representing a domain."""
+
+    admin: Any = None
+    creation_date: Any = None
+    dnssec: Any = None
+    expiration_date: Any = None
+    last_updated: Any = None
+    name_servers: Any = None
+    owner: Any = None
+    registrar: Any = None
+    reseller: Any = None
+    registrant: Any = None
+    status: Any = None
+    statuses: Any = None
+
+
+class WhoisUnknownTLD(Exception):
+    """Exception class when unknown TLD encountered."""
+
+
+class GenericWhoisError(Exception):
+    """Exception class for all other exceptions originating from the external whois library."""
+
+
+def query(domain: str) -> Domain | None:
+    """Wrap around the external whois library call and return internal domain representation."""
+
+    wh = None
+    try:
+        wh = whois.query(domain)
+    except whois.exceptions.UnknownTld as ex:
+        raise WhoisUnknownTLD from ex
+    except Exception as ex:
+        raise GenericWhoisError from ex
+    else:
+        # backward-compatibility
+        if wh is None:
+            return None
+
+        # field mapping here
+        return Domain(
+            admin=wh.admin,
+            creation_date=wh.creation_date,
+            dnssec=wh.dnssec,
+            expiration_date=wh.expiration_date,
+            last_updated=wh.last_updated,
+            name_servers=wh.name_servers,
+            owner=wh.owner,
+            registrar=wh.registrar,
+            reseller=wh.reseller,
+            registrant=wh.registrant,
+            status=wh.status,
+            statuses=wh.statuses,
+        )

homeassistant/components/whois/sensor.py

@@ -7,8 +7,6 @@
 from datetime import UTC, datetime
 from typing import cast
 
-from whois import Domain
-
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
@@ -26,6 +24,7 @@
 from homeassistant.util import dt as dt_util
 
 from .const import ATTR_EXPIRES, ATTR_NAME_SERVERS, ATTR_REGISTRAR, ATTR_UPDATED, DOMAIN
+from .helper import Domain
 
 
 @dataclass(frozen=True, kw_only=True)

homeassistant/components/whois/strings.json

@@ -8,8 +8,6 @@
       }
     },
     "error": {
-      "unexpected_response": "Unexpected response from whois server",
-      "unknown_date_format": "Unknown date format in whois server response",
       "unknown_tld": "The given TLD is unknown or not available to this integration",
       "whois_command_failed": "Whois command failed: could not retrieve whois information"
     },

tests/components/whois/conftest.py

@@ -41,12 +41,9 @@ def mock_setup_entry() -> Generator[AsyncMock, None, None]:
 @pytest.fixture
 def mock_whois() -> Generator[MagicMock, None, None]:
     """Return a mocked query."""
-    with (
-        patch(
-            "homeassistant.components.whois.whois_query",
-        ) as whois_mock,
-        patch("homeassistant.components.whois.config_flow.whois.query", new=whois_mock),
-    ):
+    with patch(
+        "homeassistant.components.whois.helper.whois.query",
+    ) as whois_mock:
         domain = whois_mock.return_value
         domain.abuse_contact = "abuse@example.com"
         domain.admin = "admin@example.com"
@@ -69,49 +66,31 @@ def mock_whois() -> Generator[MagicMock, None, None]:
 
 @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"]
-
+    """Return a mocked query that only sets one attribute."""
     with patch(
-        "homeassistant.components.whois.whois_query", LimitedWhoisMock
+        "homeassistant.components.whois.helper.whois.query",
     ) as whois_mock:
-        yield whois_mock
+        domain = whois_mock.return_value
+        domain.last_updated = datetime(
+            2022, 1, 1, 0, 0, 0, tzinfo=dt_util.get_time_zone("Europe/Amsterdam")
+        )
+        yield domain
 
 
 @pytest.fixture
-async def init_integration(
-    hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_whois: 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
+def mock_whois_empty() -> Generator[Mock, None, None]:
+    """Mock an empty response from the whois library."""
+    with patch(
+        "homeassistant.components.whois.helper.whois.query",
+    ) as whois_mock:
+        whois_mock.return_value = None
+        yield whois_mock
 
 
 @pytest.fixture
-async def init_integration_missing_some_attrs(
+async def init_integration(
     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)

tests/components/whois/snapshots/test_config_flow.ambr

@@ -1,5 +1,5 @@
 # serializer version: 1
-# name: test_full_flow_with_error[FailedParsingWhoisOutput-unexpected_response]
+# name: test_full_flow_with_error[FailedParsingWhoisOutput-whois_command_failed]
   FlowResultSnapshot({
     'context': dict({
       'source': 'user',
@@ -37,7 +37,7 @@
     'version': 1,
   })
 # ---
-# name: test_full_flow_with_error[UnknownDateFormat-unknown_date_format]
+# name: test_full_flow_with_error[UnknownDateFormat-whois_command_failed]
   FlowResultSnapshot({
     'context': dict({
       'source': 'user',

tests/components/whois/test_config_flow.py

@@ -49,8 +49,8 @@ async def test_full_user_flow(
     ("throw", "reason"),
     [
         (UnknownTld, "unknown_tld"),
-        (FailedParsingWhoisOutput, "unexpected_response"),
-        (UnknownDateFormat, "unknown_date_format"),
+        (FailedParsingWhoisOutput, "whois_command_failed"),
+        (UnknownDateFormat, "whois_command_failed"),
         (WhoisCommandFailed, "whois_command_failed"),
     ],
 )

tests/components/whois/test_diagnostics.py

@@ -1,5 +1,6 @@
 """Tests for the diagnostics data provided by the Whois integration."""
 
+import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.core import HomeAssistant
@@ -9,6 +10,7 @@
 from tests.typing import ClientSessionGenerator
 
 
+@pytest.mark.usefixtures("mock_whois")
 async def test_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,

tests/components/whois/test_init.py

@@ -3,14 +3,9 @@
 from unittest.mock import MagicMock
 
 import pytest
-from whois.exceptions import (
-    FailedParsingWhoisOutput,
-    UnknownDateFormat,
-    UnknownTld,
-    WhoisCommandFailed,
-)
 
 from homeassistant.components.whois.const import DOMAIN
+from homeassistant.components.whois.helper import GenericWhoisError, WhoisUnknownTLD
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -39,7 +34,7 @@ async def test_load_unload_config_entry(
 
 @pytest.mark.parametrize(
     "side_effect",
-    [FailedParsingWhoisOutput, UnknownDateFormat, UnknownTld, WhoisCommandFailed],
+    [WhoisUnknownTLD, GenericWhoisError],
 )
 async def test_error_handling(
     hass: HomeAssistant,

tests/components/whois/test_sensor.py

@@ -1,7 +1,5 @@
 """Tests for the sensors provided by the Whois integration."""
 
-from unittest.mock import MagicMock
-
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
@@ -14,12 +12,13 @@
 from tests.common import async_fire_time_changed
 
 pytestmark = [
-    pytest.mark.usefixtures("init_integration"),
     pytest.mark.freeze_time("2022-01-01 12:00:00", tz_offset=0),
 ]
 
 
-@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+@pytest.mark.usefixtures(
+    "entity_registry_enabled_by_default", "mock_whois", "init_integration"
+)
 @pytest.mark.parametrize(
     "entity_id",
     [
@@ -54,6 +53,7 @@ async def test_whois_sensors(
     assert device_entry == snapshot
 
 
+@pytest.mark.usefixtures("mock_whois_missing_some_attrs", "init_integration")
 async def test_whois_sensors_missing_some_attrs(
     hass: HomeAssistant, entity_registry: er.EntityRegistry, snapshot: SnapshotAssertion
 ) -> None:
@@ -65,6 +65,7 @@ async def test_whois_sensors_missing_some_attrs(
     assert entry == snapshot
 
 
+@pytest.mark.usefixtures("mock_whois", "init_integration")
 @pytest.mark.parametrize(
     "entity_id",
     [
@@ -85,7 +86,11 @@ async def test_disabled_by_default_sensors(
     assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION
 
 
-@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+@pytest.mark.usefixtures(
+    "entity_registry_enabled_by_default",
+    "mock_whois_empty",
+    "init_integration",
+)
 @pytest.mark.parametrize(
     "entity_id",
     [
@@ -100,11 +105,8 @@ async def test_disabled_by_default_sensors(
         "sensor.home_assistant_io_reseller",
     ],
 )
-async def test_no_data(
-    hass: HomeAssistant, mock_whois: MagicMock, entity_id: str
-) -> None:
+async def test_no_data(hass: HomeAssistant, entity_id: str) -> None:
     """Test whois sensors become unknown when there is no data provided."""
-    mock_whois.return_value = None
 
     async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
     await hass.async_block_till_done()