Simplify Minecraft Server SRV handling (#100726)

homeassistant/components/minecraft_server/__init__.py

@@ -4,8 +4,10 @@ from __future__ import annotations
 import logging
 from typing import Any
 
+from mcstatus import JavaServer
+
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_HOST, CONF_NAME, Platform
+from homeassistant.const import CONF_ADDRESS, CONF_HOST, CONF_PORT, Platform
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.device_registry as dr
 import homeassistant.helpers.entity_registry as er
@@ -20,20 +22,14 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Minecraft Server from a config entry."""
-    _LOGGER.debug(
-        "Creating coordinator instance for '%s' (%s)",
-        entry.data[CONF_NAME],
-        entry.data[CONF_HOST],
-    )
 
     # Create coordinator instance.
-    config_entry_id = entry.entry_id
+    coordinator = MinecraftServerCoordinator(hass, entry)
-    coordinator = MinecraftServerCoordinator(hass, config_entry_id, entry.data)
     await coordinator.async_config_entry_first_refresh()
 
     # Store coordinator instance.
     domain_data = hass.data.setdefault(DOMAIN, {})
-    domain_data[config_entry_id] = coordinator
+    domain_data[entry.entry_id] = coordinator
 
     # Set up platforms.
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@@ -43,7 +39,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload Minecraft Server config entry."""
-    config_entry_id = config_entry.entry_id
 
     # Unload platforms.
     unload_ok = await hass.config_entries.async_unload_platforms(
@@ -51,17 +46,18 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     )
 
     # Clean up.
-    hass.data[DOMAIN].pop(config_entry_id)
+    hass.data[DOMAIN].pop(config_entry.entry_id)
 
     return unload_ok
 
 
 async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Migrate old config entry to a new format."""
-    _LOGGER.debug("Migrating from version %s", config_entry.version)
 
     # 1 --> 2: Use config entry ID as base for unique IDs.
     if config_entry.version == 1:
+        _LOGGER.debug("Migrating from version 1")
+
         old_unique_id = config_entry.unique_id
         assert old_unique_id
         config_entry_id = config_entry.entry_id
@@ -78,7 +74,52 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         # Migrate entities.
         await er.async_migrate_entries(hass, config_entry_id, _migrate_entity_unique_id)
 
-    _LOGGER.debug("Migration to version %s successful", config_entry.version)
+        _LOGGER.debug("Migration to version 2 successful")
+
+    # 2 --> 3: Use address instead of host and port in config entry.
+    if config_entry.version == 2:
+        _LOGGER.debug("Migrating from version 2")
+
+        config_data = config_entry.data
+
+        # Migrate config entry.
+        try:
+            address = config_data[CONF_HOST]
+            JavaServer.lookup(address)
+            host_only_lookup_success = True
+        except ValueError as error:
+            host_only_lookup_success = False
+            _LOGGER.debug(
+                "Hostname (without port) cannot be parsed (error: %s), trying again with port",
+                error,
+            )
+
+        if not host_only_lookup_success:
+            try:
+                address = f"{config_data[CONF_HOST]}:{config_data[CONF_PORT]}"
+                JavaServer.lookup(address)
+            except ValueError as error:
+                _LOGGER.exception(
+                    "Can't migrate configuration entry due to error while parsing server address (error: %s), try again later",
+                    error,
+                )
+                return False
+
+        _LOGGER.debug(
+            "Migrating config entry, replacing host '%s' and port '%s' with address '%s'",
+            config_data[CONF_HOST],
+            config_data[CONF_PORT],
+            address,
+        )
+
+        new_data = config_data.copy()
+        new_data[CONF_ADDRESS] = address
+        del new_data[CONF_HOST]
+        del new_data[CONF_PORT]
+        config_entry.version = 3
+        hass.config_entries.async_update_entry(config_entry, data=new_data)
+
+        _LOGGER.debug("Migration to version 3 successful")
 
     return True
 

homeassistant/components/minecraft_server/config_flow.py

@@ -1,18 +1,16 @@
 """Config flow for Minecraft Server integration."""
-from contextlib import suppress
 import logging
 
 from mcstatus import JavaServer
 import voluptuous as vol
 
 from homeassistant.config_entries import ConfigFlow
-from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PORT
+from homeassistant.const import CONF_ADDRESS, CONF_NAME
 from homeassistant.data_entry_flow import FlowResult
 
-from . import helpers
+from .const import DEFAULT_NAME, DOMAIN
-from .const import DEFAULT_NAME, DEFAULT_PORT, DOMAIN
 
-DEFAULT_HOST = "localhost:25565"
+DEFAULT_ADDRESS = "localhost:25565"
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -20,51 +18,22 @@ _LOGGER = logging.getLogger(__name__)
 class MinecraftServerConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Minecraft Server."""
 
-    VERSION = 2
+    VERSION = 3
 
     async def async_step_user(self, user_input=None) -> FlowResult:
         """Handle the initial step."""
         errors = {}
 
-        if user_input is not None:
+        if user_input:
-            host = None
+            address = user_input[CONF_ADDRESS]
-            port = DEFAULT_PORT
-            title = user_input[CONF_HOST]
 
-            # Split address at last occurrence of ':'.
+            if await self._async_is_server_online(address):
-            address_left, separator, address_right = user_input[CONF_HOST].rpartition(
+                # No error was detected, create configuration entry.
-                ":"
+                config_data = {CONF_NAME: user_input[CONF_NAME], CONF_ADDRESS: address}
-            )
+                return self.async_create_entry(title=address, data=config_data)
 
-            # If no separator is found, 'rpartition' returns ('', '', original_string).
+            # Host or port invalid or server not reachable.
-            if separator == "":
+            errors["base"] = "cannot_connect"
-                host = address_right
-            else:
-                host = address_left
-                with suppress(ValueError):
-                    port = int(address_right)
-
-            # Remove '[' and ']' in case of an IPv6 address.
-            host = host.strip("[]")
-
-            # Validate port configuration (limit to user and dynamic port range).
-            if (port < 1024) or (port > 65535):
-                errors["base"] = "invalid_port"
-            # Validate host and port by checking the server connection.
-            else:
-                # Create server instance with configuration data and ping the server.
-                config_data = {
-                    CONF_NAME: user_input[CONF_NAME],
-                    CONF_HOST: host,
-                    CONF_PORT: port,
-                }
-                if await self._async_is_server_online(host, port):
-                    # Configuration data are available and no error was detected,
-                    # create configuration entry.
-                    return self.async_create_entry(title=title, data=config_data)
-
-                # Host or port invalid or server not reachable.
-                errors["base"] = "cannot_connect"
 
         # Show configuration form (default form in case of no user_input,
         # form filled with user_input and eventually with errors otherwise).
@@ -83,24 +52,32 @@ class MinecraftServerConfigFlow(ConfigFlow, domain=DOMAIN):
                         CONF_NAME, default=user_input.get(CONF_NAME, DEFAULT_NAME)
                     ): str,
                     vol.Required(
-                        CONF_HOST, default=user_input.get(CONF_HOST, DEFAULT_HOST)
+                        CONF_ADDRESS,
+                        default=user_input.get(CONF_ADDRESS, DEFAULT_ADDRESS),
                     ): vol.All(str, vol.Lower),
                 }
             ),
             errors=errors,
         )
 
-    async def _async_is_server_online(self, host: str, port: int) -> bool:
+    async def _async_is_server_online(self, address: str) -> bool:
         """Check server connection using a 'status' request and return result."""
 
-        # Check if host is a SRV record. If so, update server data.
+        # Parse and check server address.
-        if srv_record := await helpers.async_check_srv_record(host):
+        try:
-            # Use extracted host and port from SRV record.
+            server = await JavaServer.async_lookup(address)
-            host = srv_record[CONF_HOST]
+        except ValueError as error:
-            port = srv_record[CONF_PORT]
+            _LOGGER.debug(
+                (
+                    "Error occurred while parsing server address '%s' -"
+                    " ValueError: %s"
+                ),
+                address,
+                error,
+            )
+            return False
 
         # Send a status request to the server.
-        server = JavaServer(host, port)
         try:
             await server.async_status()
             return True
@@ -110,8 +87,8 @@ class MinecraftServerConfigFlow(ConfigFlow, domain=DOMAIN):
                     "Error occurred while trying to check the connection to '%s:%s' -"
                     " OSError: %s"
                 ),
-                host,
+                server.address.host,
-                port,
+                server.address.port,
                 error,
             )
 

homeassistant/components/minecraft_server/const.py

@@ -1,7 +1,6 @@
 """Constants for the Minecraft Server integration."""
 
 DEFAULT_NAME = "Minecraft Server"
-DEFAULT_PORT = 25565
 
 DOMAIN = "minecraft_server"
 

homeassistant/components/minecraft_server/coordinator.py

@@ -1,20 +1,18 @@
 """The Minecraft Server integration."""
 from __future__ import annotations
 
-from collections.abc import Mapping
 from dataclasses import dataclass
 from datetime import timedelta
 import logging
-from typing import Any
 
 from mcstatus.server import JavaServer
 
-from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PORT
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_ADDRESS, CONF_NAME
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from . import helpers
-
 SCAN_INTERVAL = timedelta(seconds=60)
 
 _LOGGER = logging.getLogger(__name__)
@@ -36,12 +34,11 @@ class MinecraftServerData:
 class MinecraftServerCoordinator(DataUpdateCoordinator[MinecraftServerData]):
     """Minecraft Server data update coordinator."""
 
-    _srv_record_checked = False
+    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
-
-    def __init__(
-        self, hass: HomeAssistant, unique_id: str, config_data: Mapping[str, Any]
-    ) -> None:
         """Initialize coordinator instance."""
+        config_data = config_entry.data
+        self.unique_id = config_entry.entry_id
+
         super().__init__(
             hass=hass,
             name=config_data[CONF_NAME],
@@ -49,34 +46,20 @@ class MinecraftServerCoordinator(DataUpdateCoordinator[MinecraftServerData]):
             update_interval=SCAN_INTERVAL,
         )
 
-        # Server data
+        try:
-        self.unique_id = unique_id
+            self._server = JavaServer.lookup(config_data[CONF_ADDRESS])
-        self._host = config_data[CONF_HOST]
+        except ValueError as error:
-        self._port = config_data[CONF_PORT]
+            raise HomeAssistantError(
-
+                f"Address in configuration entry cannot be parsed (error: {error}), please remove this device and add it again"
-        # 3rd party library instance
+            ) from error
-        self._server = JavaServer(self._host, self._port)
 
     async def _async_update_data(self) -> MinecraftServerData:
         """Get server data from 3rd party library and update properties."""
-
-        # Check once if host is a valid Minecraft SRV record.
-        if not self._srv_record_checked:
-            self._srv_record_checked = True
-            if srv_record := await helpers.async_check_srv_record(self._host):
-                # Overwrite host, port and 3rd party library instance
-                # with data extracted out of the SRV record.
-                self._host = srv_record[CONF_HOST]
-                self._port = srv_record[CONF_PORT]
-                self._server = JavaServer(self._host, self._port)
-
-        # Send status request to the server.
         try:
             status_response = await self._server.async_status()
         except OSError as error:
             raise UpdateFailed(error) from error
 
-        # Got answer to request, update properties.
         players_list = []
         if players := status_response.players.sample:
             for player in players:

homeassistant/components/minecraft_server/helpers.py

@@ -1,38 +0,0 @@
-"""Helper functions of Minecraft Server integration."""
-import logging
-from typing import Any
-
-import aiodns
-
-from homeassistant.const import CONF_HOST, CONF_PORT
-
-SRV_RECORD_PREFIX = "_minecraft._tcp"
-
-_LOGGER = logging.getLogger(__name__)
-
-
-async def async_check_srv_record(host: str) -> dict[str, Any] | None:
-    """Check if the given host is a valid Minecraft SRV record."""
-    srv_record = None
-
-    try:
-        srv_query = await aiodns.DNSResolver().query(
-            host=f"{SRV_RECORD_PREFIX}.{host}", qtype="SRV"
-        )
-    except aiodns.error.DNSError:
-        # 'host' is not a Minecraft SRV record.
-        pass
-    else:
-        # 'host' is a valid Minecraft SRV record, extract the data.
-        srv_record = {
-            CONF_HOST: srv_query[0].host,
-            CONF_PORT: srv_query[0].port,
-        }
-        _LOGGER.debug(
-            "'%s' is a valid Minecraft SRV record ('%s:%s')",
-            host,
-            srv_record[CONF_HOST],
-            srv_record[CONF_PORT],
-        )
-
-    return srv_record

homeassistant/components/minecraft_server/manifest.json

@@ -7,5 +7,5 @@
   "iot_class": "local_polling",
   "loggers": ["dnspython", "mcstatus"],
   "quality_scale": "silver",
-  "requirements": ["aiodns==3.0.0", "mcstatus==11.0.0"]
+  "requirements": ["mcstatus==11.0.0"]
 }

homeassistant/components/minecraft_server/strings.json

@@ -6,13 +6,12 @@
         "description": "Set up your Minecraft Server instance to allow monitoring.",
         "data": {
           "name": "[%key:common::config_flow::data::name%]",
-          "host": "[%key:common::config_flow::data::host%]"
+          "address": "Server address"
         }
       }
     },
     "error": {
-      "invalid_port": "Port must be in range from 1024 to 65535. Please correct it and try again.",
+      "cannot_connect": "Failed to connect to server. Please check the address and try again. If a port was provided, it must be within a valid range. Also ensure that you are running at least version 1.7 of Minecraft Java Edition on your server."
-      "cannot_connect": "Failed to connect to server. Please check the host and port and try again. Also ensure that you are running at least Minecraft version 1.7 on your server."
     }
   },
   "entity": {

requirements_all.txt

@@ -216,7 +216,6 @@ aiocomelit==0.0.8
 aiodiscover==1.5.1
 
 # homeassistant.components.dnsip
-# homeassistant.components.minecraft_server
 aiodns==3.0.0
 
 # homeassistant.components.eafm

requirements_test_all.txt

@@ -197,7 +197,6 @@ aiocomelit==0.0.8
 aiodiscover==1.5.1
 
 # homeassistant.components.dnsip
-# homeassistant.components.minecraft_server
 aiodns==3.0.0
 
 # homeassistant.components.eafm

tests/components/minecraft_server/const.py

@@ -7,6 +7,8 @@ from mcstatus.status_response import (
 )
 
 TEST_HOST = "mc.dummyserver.com"
+TEST_PORT = 25566
+TEST_ADDRESS = f"{TEST_HOST}:{TEST_PORT}"
 
 TEST_JAVA_STATUS_RESPONSE_RAW = {
     "description": {"text": "Dummy Description"},

tests/components/minecraft_server/test_config_flow.py

@@ -1,59 +1,20 @@
 """Tests for the Minecraft Server config flow."""
 
-from unittest.mock import AsyncMock, patch
+from unittest.mock import patch
 
-import aiodns
+from mcstatus import JavaServer
 
-from homeassistant.components.minecraft_server.const import (
+from homeassistant.components.minecraft_server.const import DEFAULT_NAME, DOMAIN
-    DEFAULT_NAME,
-    DEFAULT_PORT,
-    DOMAIN,
-)
 from homeassistant.config_entries import SOURCE_USER
-from homeassistant.const import CONF_HOST, CONF_NAME
+from homeassistant.const import CONF_ADDRESS, CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
-from .const import TEST_HOST, TEST_JAVA_STATUS_RESPONSE
+from .const import TEST_ADDRESS, TEST_HOST, TEST_JAVA_STATUS_RESPONSE, TEST_PORT
-
-
-class QueryMock:
-    """Mock for result of aiodns.DNSResolver.query."""
-
-    def __init__(self) -> None:
-        """Set up query result mock."""
-        self.host = TEST_HOST
-        self.port = 23456
-        self.priority = 1
-        self.weight = 1
-        self.ttl = None
-
 
 USER_INPUT = {
     CONF_NAME: DEFAULT_NAME,
-    CONF_HOST: f"{TEST_HOST}:{DEFAULT_PORT}",
+    CONF_ADDRESS: TEST_ADDRESS,
-}
-
-USER_INPUT_SRV = {CONF_NAME: DEFAULT_NAME, CONF_HOST: TEST_HOST}
-
-USER_INPUT_IPV4 = {
-    CONF_NAME: DEFAULT_NAME,
-    CONF_HOST: f"1.1.1.1:{DEFAULT_PORT}",
-}
-
-USER_INPUT_IPV6 = {
-    CONF_NAME: DEFAULT_NAME,
-    CONF_HOST: f"[::ffff:0101:0101]:{DEFAULT_PORT}",
-}
-
-USER_INPUT_PORT_TOO_SMALL = {
-    CONF_NAME: DEFAULT_NAME,
-    CONF_HOST: f"{TEST_HOST}:1023",
-}
-
-USER_INPUT_PORT_TOO_LARGE = {
-    CONF_NAME: DEFAULT_NAME,
-    CONF_HOST: f"{TEST_HOST}:65536",
 }
 
 
@@ -67,39 +28,25 @@ async def test_show_config_form(hass: HomeAssistant) -> None:
     assert result["step_id"] == "user"
 
 
-async def test_port_too_small(hass: HomeAssistant) -> None:
+async def test_lookup_failed(hass: HomeAssistant) -> None:
-    """Test error in case of a too small port."""
+    """Test error in case of a failed connection."""
     with patch(
-        "aiodns.DNSResolver.query",
+        "mcstatus.server.JavaServer.async_lookup",
-        side_effect=aiodns.error.DNSError,
+        side_effect=ValueError,
     ):
         result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT_PORT_TOO_SMALL
+            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT
         )
 
         assert result["type"] == FlowResultType.FORM
-        assert result["errors"] == {"base": "invalid_port"}
+        assert result["errors"] == {"base": "cannot_connect"}
-
-
-async def test_port_too_large(hass: HomeAssistant) -> None:
-    """Test error in case of a too large port."""
-    with patch(
-        "aiodns.DNSResolver.query",
-        side_effect=aiodns.error.DNSError,
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT_PORT_TOO_LARGE
-        )
-
-        assert result["type"] == FlowResultType.FORM
-        assert result["errors"] == {"base": "invalid_port"}
 
 
 async def test_connection_failed(hass: HomeAssistant) -> None:
     """Test error in case of a failed connection."""
     with patch(
-        "aiodns.DNSResolver.query",
+        "mcstatus.server.JavaServer.async_lookup",
-        side_effect=aiodns.error.DNSError,
+        return_value=JavaServer(host=TEST_HOST, port=TEST_PORT),
     ), patch("mcstatus.server.JavaServer.async_status", side_effect=OSError):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT
@@ -109,30 +56,11 @@ async def test_connection_failed(hass: HomeAssistant) -> None:
         assert result["errors"] == {"base": "cannot_connect"}
 
 
-async def test_connection_succeeded_with_srv_record(hass: HomeAssistant) -> None:
+async def test_connection_succeeded(hass: HomeAssistant) -> None:
-    """Test config entry in case of a successful connection with a SRV record."""
-    with patch(
-        "aiodns.DNSResolver.query",
-        side_effect=AsyncMock(return_value=[QueryMock()]),
-    ), patch(
-        "mcstatus.server.JavaServer.async_status",
-        return_value=TEST_JAVA_STATUS_RESPONSE,
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT_SRV
-        )
-
-        assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["title"] == USER_INPUT_SRV[CONF_HOST]
-        assert result["data"][CONF_NAME] == USER_INPUT_SRV[CONF_NAME]
-        assert result["data"][CONF_HOST] == USER_INPUT_SRV[CONF_HOST]
-
-
-async def test_connection_succeeded_with_host(hass: HomeAssistant) -> None:
     """Test config entry in case of a successful connection with a host name."""
     with patch(
-        "aiodns.DNSResolver.query",
+        "mcstatus.server.JavaServer.async_lookup",
-        side_effect=aiodns.error.DNSError,
+        return_value=JavaServer(host=TEST_HOST, port=TEST_PORT),
     ), patch(
         "mcstatus.server.JavaServer.async_status",
         return_value=TEST_JAVA_STATUS_RESPONSE,
@@ -142,44 +70,6 @@ async def test_connection_succeeded_with_host(hass: HomeAssistant) -> None:
         )
 
         assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["title"] == USER_INPUT[CONF_HOST]
+        assert result["title"] == USER_INPUT[CONF_ADDRESS]
         assert result["data"][CONF_NAME] == USER_INPUT[CONF_NAME]
-        assert result["data"][CONF_HOST] == TEST_HOST
+        assert result["data"][CONF_ADDRESS] == TEST_ADDRESS
-
-
-async def test_connection_succeeded_with_ip4(hass: HomeAssistant) -> None:
-    """Test config entry in case of a successful connection with an IPv4 address."""
-    with patch(
-        "aiodns.DNSResolver.query",
-        side_effect=aiodns.error.DNSError,
-    ), patch(
-        "mcstatus.server.JavaServer.async_status",
-        return_value=TEST_JAVA_STATUS_RESPONSE,
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT_IPV4
-        )
-
-        assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["title"] == USER_INPUT_IPV4[CONF_HOST]
-        assert result["data"][CONF_NAME] == USER_INPUT_IPV4[CONF_NAME]
-        assert result["data"][CONF_HOST] == "1.1.1.1"
-
-
-async def test_connection_succeeded_with_ip6(hass: HomeAssistant) -> None:
-    """Test config entry in case of a successful connection with an IPv6 address."""
-    with patch(
-        "aiodns.DNSResolver.query",
-        side_effect=aiodns.error.DNSError,
-    ), patch(
-        "mcstatus.server.JavaServer.async_status",
-        return_value=TEST_JAVA_STATUS_RESPONSE,
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=USER_INPUT_IPV6
-        )
-
-        assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert result["title"] == USER_INPUT_IPV6[CONF_HOST]
-        assert result["data"][CONF_NAME] == USER_INPUT_IPV6[CONF_NAME]
-        assert result["data"][CONF_HOST] == "::ffff:0101:0101"

tests/components/minecraft_server/test_init.py

@@ -1,24 +1,20 @@
 """Tests for the Minecraft Server integration."""
 from unittest.mock import patch
 
-import aiodns
+from mcstatus import JavaServer
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
-from homeassistant.components.minecraft_server.const import (
+from homeassistant.components.minecraft_server.const import DEFAULT_NAME, DOMAIN
-    DEFAULT_NAME,
-    DEFAULT_PORT,
-    DOMAIN,
-)
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
-from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PORT
+from homeassistant.const import CONF_ADDRESS, CONF_HOST, CONF_NAME, CONF_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from .const import TEST_HOST, TEST_JAVA_STATUS_RESPONSE
+from .const import TEST_ADDRESS, TEST_HOST, TEST_JAVA_STATUS_RESPONSE, TEST_PORT
 
 from tests.common import MockConfigEntry
 
-TEST_UNIQUE_ID = f"{TEST_HOST}-{DEFAULT_PORT}"
+TEST_UNIQUE_ID = f"{TEST_HOST}-{TEST_PORT}"
 
 SENSOR_KEYS = [
     {"v1": "Latency Time", "v2": "latency"},
@@ -32,43 +28,54 @@ SENSOR_KEYS = [
 BINARY_SENSOR_KEYS = {"v1": "Status", "v2": "status"}
 
 
-async def test_entry_migration_v1_to_v2(hass: HomeAssistant) -> None:
+def create_v1_mock_config_entry(hass: HomeAssistant) -> int:
-    """Test entry migration from version 1 to 2."""
+    """Create mock config entry."""
-
-    # Create mock config entry.
     config_entry_v1 = MockConfigEntry(
         domain=DOMAIN,
         unique_id=TEST_UNIQUE_ID,
         data={
             CONF_NAME: DEFAULT_NAME,
             CONF_HOST: TEST_HOST,
-            CONF_PORT: DEFAULT_PORT,
+            CONF_PORT: TEST_PORT,
         },
         version=1,
     )
     config_entry_id = config_entry_v1.entry_id
     config_entry_v1.add_to_hass(hass)
 
-    # Create mock device entry.
+    return config_entry_id
+
+
+def create_v1_mock_device_entry(hass: HomeAssistant, config_entry_id: int) -> int:
+    """Create mock device entry."""
     device_registry = dr.async_get(hass)
     device_entry_v1 = device_registry.async_get_or_create(
         config_entry_id=config_entry_id,
         identifiers={(DOMAIN, TEST_UNIQUE_ID)},
     )
     device_entry_id = device_entry_v1.id
+
     assert device_entry_v1
     assert device_entry_id
 
-    # Create mock sensor entity entries.
+    return device_entry_id
+
+
+def create_v1_mock_sensor_entity_entries(
+    hass: HomeAssistant, config_entry_id: int, device_entry_id: int
+) -> list[dict]:
+    """Create mock sensor entity entries."""
     sensor_entity_id_key_mapping_list = []
+    config_entry = hass.config_entries.async_get_entry(config_entry_id)
     entity_registry = er.async_get(hass)
+
     for sensor_key in SENSOR_KEYS:
         entity_unique_id = f"{TEST_UNIQUE_ID}-{sensor_key['v1']}"
         entity_entry_v1 = entity_registry.async_get_or_create(
             SENSOR_DOMAIN,
             DOMAIN,
             unique_id=entity_unique_id,
-            config_entry=config_entry_v1,
+            config_entry=config_entry,
             device_id=device_entry_id,
         )
         assert entity_entry_v1.unique_id == entity_unique_id
@@ -76,25 +83,51 @@ async def test_entry_migration_v1_to_v2(hass: HomeAssistant) -> None:
             {"entity_id": entity_entry_v1.entity_id, "key": sensor_key["v2"]}
         )
 
-    # Create mock binary sensor entity entry.
+    return sensor_entity_id_key_mapping_list
+
+
+def create_v1_mock_binary_sensor_entity_entry(
+    hass: HomeAssistant, config_entry_id: int, device_entry_id: int
+) -> dict:
+    """Create mock binary sensor entity entry."""
+    config_entry = hass.config_entries.async_get_entry(config_entry_id)
+    entity_registry = er.async_get(hass)
     entity_unique_id = f"{TEST_UNIQUE_ID}-{BINARY_SENSOR_KEYS['v1']}"
-    entity_entry_v1 = entity_registry.async_get_or_create(
+    entity_entry = entity_registry.async_get_or_create(
         BINARY_SENSOR_DOMAIN,
         DOMAIN,
         unique_id=entity_unique_id,
-        config_entry=config_entry_v1,
+        config_entry=config_entry,
         device_id=device_entry_id,
     )
-    assert entity_entry_v1.unique_id == entity_unique_id
+    assert entity_entry.unique_id == entity_unique_id
     binary_sensor_entity_id_key_mapping = {
-        "entity_id": entity_entry_v1.entity_id,
+        "entity_id": entity_entry.entity_id,
         "key": BINARY_SENSOR_KEYS["v2"],
     }
 
+    return binary_sensor_entity_id_key_mapping
+
+
+async def test_entry_migration(hass: HomeAssistant) -> None:
+    """Test entry migration from version 1 to 3, where host and port is required for the connection to the server."""
+    config_entry_id = create_v1_mock_config_entry(hass)
+    device_entry_id = create_v1_mock_device_entry(hass, config_entry_id)
+    sensor_entity_id_key_mapping_list = create_v1_mock_sensor_entity_entries(
+        hass, config_entry_id, device_entry_id
+    )
+    binary_sensor_entity_id_key_mapping = create_v1_mock_binary_sensor_entity_entry(
+        hass, config_entry_id, device_entry_id
+    )
+
     # Trigger migration.
     with patch(
-        "aiodns.DNSResolver.query",
+        "mcstatus.server.JavaServer.lookup",
-        side_effect=aiodns.error.DNSError,
+        side_effect=[
+            ValueError,
+            JavaServer(host=TEST_HOST, port=TEST_PORT),
+            JavaServer(host=TEST_HOST, port=TEST_PORT),
+        ],
     ), patch(
         "mcstatus.server.JavaServer.async_status",
         return_value=TEST_JAVA_STATUS_RESPONSE,
@@ -103,29 +136,84 @@ async def test_entry_migration_v1_to_v2(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     # Test migrated config entry.
-    config_entry_v2 = hass.config_entries.async_get_entry(config_entry_id)
+    config_entry = hass.config_entries.async_get_entry(config_entry_id)
-    assert config_entry_v2.unique_id is None
+    assert config_entry.unique_id is None
-    assert config_entry_v2.data == {
+    assert config_entry.data == {
         CONF_NAME: DEFAULT_NAME,
-        CONF_HOST: TEST_HOST,
+        CONF_ADDRESS: TEST_ADDRESS,
-        CONF_PORT: DEFAULT_PORT,
     }
-    assert config_entry_v2.version == 2
+    assert config_entry.version == 3
 
     # Test migrated device entry.
-    device_entry_v2 = device_registry.async_get(device_entry_id)
+    device_registry = dr.async_get(hass)
-    assert device_entry_v2.identifiers == {(DOMAIN, config_entry_id)}
+    device_entry = device_registry.async_get(device_entry_id)
+    assert device_entry.identifiers == {(DOMAIN, config_entry_id)}
 
     # Test migrated sensor entity entries.
+    entity_registry = er.async_get(hass)
     for mapping in sensor_entity_id_key_mapping_list:
-        entity_entry_v2 = entity_registry.async_get(mapping["entity_id"])
+        entity_entry = entity_registry.async_get(mapping["entity_id"])
-        assert entity_entry_v2.unique_id == f"{config_entry_id}-{mapping['key']}"
+        assert entity_entry.unique_id == f"{config_entry_id}-{mapping['key']}"
 
     # Test migrated binary sensor entity entry.
-    entity_entry_v2 = entity_registry.async_get(
+    entity_entry = entity_registry.async_get(
         binary_sensor_entity_id_key_mapping["entity_id"]
     )
     assert (
-        entity_entry_v2.unique_id
+        entity_entry.unique_id
         == f"{config_entry_id}-{binary_sensor_entity_id_key_mapping['key']}"
     )
+
+
+async def test_entry_migration_host_only(hass: HomeAssistant) -> None:
+    """Test entry migration from version 1 to 3, where host alone is sufficient for the connection to the server."""
+    config_entry_id = create_v1_mock_config_entry(hass)
+    device_entry_id = create_v1_mock_device_entry(hass, config_entry_id)
+    create_v1_mock_sensor_entity_entries(hass, config_entry_id, device_entry_id)
+    create_v1_mock_binary_sensor_entity_entry(hass, config_entry_id, device_entry_id)
+
+    # Trigger migration.
+    with patch(
+        "mcstatus.server.JavaServer.lookup",
+        side_effect=[
+            JavaServer(host=TEST_HOST, port=TEST_PORT),
+            JavaServer(host=TEST_HOST, port=TEST_PORT),
+        ],
+    ), patch(
+        "mcstatus.server.JavaServer.async_status",
+        return_value=TEST_JAVA_STATUS_RESPONSE,
+    ):
+        assert await hass.config_entries.async_setup(config_entry_id)
+        await hass.async_block_till_done()
+
+    # Test migrated config entry.
+    config_entry = hass.config_entries.async_get_entry(config_entry_id)
+    assert config_entry.unique_id is None
+    assert config_entry.data == {
+        CONF_NAME: DEFAULT_NAME,
+        CONF_ADDRESS: TEST_HOST,
+    }
+    assert config_entry.version == 3
+
+
+async def test_entry_migration_v3_failure(hass: HomeAssistant) -> None:
+    """Test failed entry migration from version 2 to 3."""
+    config_entry_id = create_v1_mock_config_entry(hass)
+    device_entry_id = create_v1_mock_device_entry(hass, config_entry_id)
+    create_v1_mock_sensor_entity_entries(hass, config_entry_id, device_entry_id)
+    create_v1_mock_binary_sensor_entity_entry(hass, config_entry_id, device_entry_id)
+
+    # Trigger migration.
+    with patch(
+        "mcstatus.server.JavaServer.lookup",
+        side_effect=[
+            ValueError,
+            ValueError,
+        ],
+    ):
+        assert not await hass.config_entries.async_setup(config_entry_id)
+        await hass.async_block_till_done()
+
+    # Test config entry.
+    config_entry = hass.config_entries.async_get_entry(config_entry_id)
+    assert config_entry.version == 2