home-assistant · balloob · May 23, 2021 · May 21, 2021
diff --git a/homeassistant/components/ps4/__init__.py b/homeassistant/components/ps4/__init__.py
@@ -25,7 +25,14 @@
 from homeassistant.util.json import load_json, save_json
 
 from .config_flow import PlayStation4FlowHandler  # noqa: F401
-from .const import ATTR_MEDIA_IMAGE_URL, COMMANDS, DOMAIN, GAMES_FILE, PS4_DATA
+from .const import (
+    ATTR_MEDIA_IMAGE_URL,
+    COMMANDS,
+    COUNTRYCODE_NAMES,
+    DOMAIN,
+    GAMES_FILE,
+    PS4_DATA,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -91,7 +98,7 @@ async def async_migrate_entry(hass, entry):
             hass.helpers.aiohttp_client.async_get_clientsession()
         )
         if loc:
-            country = loc.country_name
+            country = COUNTRYCODE_NAMES.get(loc.country_code)
             if country in COUNTRIES:
                 for device in data["devices"]:
                     device[CONF_REGION] = country

diff --git a/homeassistant/components/ps4/config_flow.py b/homeassistant/components/ps4/config_flow.py
@@ -17,7 +17,13 @@
 )
 from homeassistant.util import location
 
-from .const import CONFIG_ENTRY_VERSION, DEFAULT_ALIAS, DEFAULT_NAME, DOMAIN
+from .const import (
+    CONFIG_ENTRY_VERSION,
+    COUNTRYCODE_NAMES,
+    DEFAULT_ALIAS,
+    DEFAULT_NAME,
+    DOMAIN,
+)
 
 CONF_MODE = "Config Mode"
 CONF_AUTO = "Auto Discover"
@@ -178,7 +184,7 @@ async def async_step_link(self, user_input=None):
                 self.hass.helpers.aiohttp_client.async_get_clientsession()
             )
         if self.location:
-            country = self.location.country_name
+            country = COUNTRYCODE_NAMES.get(self.location.country_code)
             if country in COUNTRIES:
                 default_region = country
 

diff --git a/homeassistant/components/ps4/const.py b/homeassistant/components/ps4/const.py
@@ -12,3 +12,71 @@
 
 # Deprecated used for logger/backwards compatibility from 0.89
 REGIONS = ["R1", "R2", "R3", "R4", "R5"]
+
+COUNTRYCODE_NAMES = {
+    "AE": "United Arab Emirates",
+    "AR": "Argentina",
+    "AT": "Austria",
+    "AU": "Australia",
+    "BE": "Belgium",
+    "BG": "Bulgaria",
+    "BH": "Bahrain",
+    "BR": "Brazil",
+    "CA": "Canada",
+    "CH": "Switzerland",
+    "CL": "Chile",
+    "CO": "Columbia",
+    "CR": "Costa Rica",
+    "CY": "Cyprus",
+    "CZ": "Czech Republic",
+    "DE": "Germany",
+    "DK": "Denmark",
+    "EC": "Ecuador",
+    "ES": "Spain",
+    "FI": "Finland",
+    "FR": "France",
+    "GB": "United Kingdom",
+    "GR": "Greece",
+    "GT": "Guatemala",
+    "HK": "Hong Kong",
+    "HN": "Honduras",
+    "HR": "Croatia",
+    "HU": "Hungary",
+    "ID": "Indonesia",
+    "IE": "Ireland",
+    "IL": "Israel",
+    "IN": "India",
+    "IS": "Iceland",
+    "IT": "Italy",
+    "JP": "Japan",
+    "KW": "Kuwait",
+    "LB": "Lebanon",
+    "LU": "Luxembourg",
+    "MT": "Malta",
+    "MX": "Mexico",
+    "MY": "Maylasia",
+    "NI": "Nicaragua",
+    "NL": "Nederland",
+    "NO": "Norway",
+    "NZ": "New Zealand",
+    "OM": "Oman",
+    "PA": "Panama",
+    "PE": "Peru",
+    "PL": "Poland",
+    "PT": "Portugal",
+    "QA": "Qatar",
+    "RO": "Romania",
+    "RU": "Russia",
+    "SA": "Saudi Arabia",
+    "SE": "Sweden",
+    "SG": "Singapore",
+    "SI": "Slovenia",
+    "SK": "Slovakia",
+    "SV": "El Salvador",
+    "TH": "Thailand",
+    "TR": "Turkey",
+    "TW": "Taiwan",
+    "UA": "Ukraine",
+    "US": "United States",
+    "ZA": "South Africa",
+}
@@ -12,9 +12,7 @@
 
 import aiohttp
 
-ELEVATION_URL = "https://api.open-elevation.com/api/v1/lookup"
-IP_API = "http://ip-api.com/json"
-IPAPI = "https://ipapi.co/json/"
+WHOAMI_URL = "https://whoami.home-assistant.io/v1"
 
 # Constants from https://github.com/maurycyp/vincenty
 # Earth ellipsoid according to WGS 84
@@ -34,7 +32,6 @@
     [
         "ip",
         "country_code",
-        "country_name",
         "region_code",
         "region_name",
         "city",
@@ -51,10 +48,7 @@ async def async_detect_location_info(
     session: aiohttp.ClientSession,
 ) -> LocationInfo | None:
     """Detect location information."""
-    data = await _get_ipapi(session)
-
-    if data is None:
-        data = await _get_ip_api(session)
+    data = await _get_whoami(session)
 
     if data is None:
         return None
@@ -164,10 +158,10 @@ def vincenty(
     return round(s, 6)
 
 
-async def _get_ipapi(session: aiohttp.ClientSession) -> dict[str, Any] | None:
-    """Query ipapi.co for location data."""
+async def _get_whoami(session: aiohttp.ClientSession) -> dict[str, Any] | None:
+    """Query whoami.home-assistant.io for location data."""
     try:
-        resp = await session.get(IPAPI, timeout=5)
+        resp = await session.get(WHOAMI_URL, timeout=30)
     except (aiohttp.ClientError, asyncio.TimeoutError):
         return None
 
@@ -176,44 +170,14 @@ async def _get_ipapi(session: aiohttp.ClientSession) -> dict[str, Any] | None:
     except (aiohttp.ClientError, ValueError):
         return None
 
-    # ipapi allows 30k free requests/month. Some users exhaust those.
-    if raw_info.get("latitude") == "Sign up to access":
-        return None
-
     return {
         "ip": raw_info.get("ip"),
         "country_code": raw_info.get("country"),
-        "country_name": raw_info.get("country_name"),
         "region_code": raw_info.get("region_code"),
         "region_name": raw_info.get("region"),
         "city": raw_info.get("city"),
-        "zip_code": raw_info.get("postal"),
-        "time_zone": raw_info.get("timezone"),
-        "latitude": raw_info.get("latitude"),
-        "longitude": raw_info.get("longitude"),
-    }
-
-
-async def _get_ip_api(session: aiohttp.ClientSession) -> dict[str, Any] | None:
-    """Query ip-api.com for location data."""
-    try:
-        resp = await session.get(IP_API, timeout=5)
-    except (aiohttp.ClientError, asyncio.TimeoutError):
-        return None
-
-    try:
-        raw_info = await resp.json()
-    except (aiohttp.ClientError, ValueError):
-        return None
-    return {
-        "ip": raw_info.get("query"),
-        "country_code": raw_info.get("countryCode"),
-        "country_name": raw_info.get("country"),
-        "region_code": raw_info.get("region"),
-        "region_name": raw_info.get("regionName"),
-        "city": raw_info.get("city"),
-        "zip_code": raw_info.get("zip"),
+        "zip_code": raw_info.get("postal_code"),
         "time_zone": raw_info.get("timezone"),
-        "latitude": raw_info.get("lat"),
-        "longitude": raw_info.get("lon"),
+        "latitude": float(raw_info.get("latitude")),
+        "longitude": float(raw_info.get("longitude")),
     }
@@ -144,7 +144,6 @@ async def test_detect_config_fail(hass, client):
         return_value=location.LocationInfo(
             ip=None,
             country_code=None,
-            country_name=None,
             region_code=None,
             region_name=None,
             city=None,

diff --git a/tests/components/ps4/test_config_flow.py b/tests/components/ps4/test_config_flow.py
@@ -64,7 +64,6 @@
 MOCK_LOCATION = location.LocationInfo(
     "0.0.0.0",
     "US",
-    "United States",
     "CA",
     "California",
     "San Diego",

diff --git a/tests/components/ps4/test_init.py b/tests/components/ps4/test_init.py
@@ -56,7 +56,6 @@
 MOCK_LOCATION = location.LocationInfo(
     "0.0.0.0",
     "US",
-    "United States",
     "CA",
     "California",
     "San Diego",

diff --git a/tests/fixtures/ip-api.com.json b/tests/fixtures/ip-api.com.json
diff --git a/tests/fixtures/ipapi.co.json b/tests/fixtures/ipapi.co.json
diff --git a/tests/fixtures/whoami.json b/tests/fixtures/whoami.json
@@ -0,0 +1,14 @@
+{
+    "ip": "1.2.3.4",
+    "city": "Gotham",
+    "continent": "Earth",
+    "country": "XX",
+    "latitude": "12.34567",
+    "longitude": "12.34567",
+    "postal_code": "12345",
+    "region_code": "00",
+    "region": "Gotham",
+    "timezone": "Earth/Gotham",
+    "iso_time": "2021-05-12T11:29:15.752Z",
+    "timestamp": 1620818956
+}
diff --git a/tests/util/test_location.py b/tests/util/test_location.py
@@ -1,5 +1,5 @@
 """Test Home Assistant location util methods."""
-from unittest.mock import Mock, patch
+from unittest.mock import Mock
 
 import aiohttp
 import pytest
@@ -72,76 +72,25 @@ def test_get_miles():
     assert round(miles, 2) == DISTANCE_MILES
 
 
-async def test_detect_location_info_ipapi(aioclient_mock, session):
-    """Test detect location info using ipapi.co."""
-    aioclient_mock.get(location_util.IPAPI, text=load_fixture("ipapi.co.json"))
+async def test_detect_location_info_whoami(aioclient_mock, session):
+    """Test detect location info using whoami.home-assistant.io."""
+    aioclient_mock.get(location_util.WHOAMI_URL, text=load_fixture("whoami.json"))
 
     info = await location_util.async_detect_location_info(session, _test_real=True)
 
     assert info is not None
     assert info.ip == "1.2.3.4"
-    assert info.country_code == "CH"
-    assert info.country_name == "Switzerland"
-    assert info.region_code == "BE"
-    assert info.region_name == "Bern"
-    assert info.city == "Bern"
-    assert info.zip_code == "3000"
-    assert info.time_zone == "Europe/Zurich"
-    assert info.latitude == 46.9480278
-    assert info.longitude == 7.4490812
+    assert info.country_code == "XX"
+    assert info.region_code == "00"
+    assert info.city == "Gotham"
+    assert info.zip_code == "12345"
+    assert info.time_zone == "Earth/Gotham"
+    assert info.latitude == 12.34567
+    assert info.longitude == 12.34567
     assert info.use_metric
 
 
-async def test_detect_location_info_ipapi_exhaust(aioclient_mock, session):
-    """Test detect location info using ipapi.co."""
-    aioclient_mock.get(location_util.IPAPI, json={"latitude": "Sign up to access"})
-    aioclient_mock.get(location_util.IP_API, text=load_fixture("ip-api.com.json"))
-
-    info = await location_util.async_detect_location_info(session, _test_real=True)
-
-    assert info is not None
-    # ip_api result because ipapi got skipped
-    assert info.country_code == "US"
-    assert len(aioclient_mock.mock_calls) == 2
-
-
-async def test_detect_location_info_ip_api(aioclient_mock, session):
-    """Test detect location info using ip-api.com."""
-    aioclient_mock.get(location_util.IP_API, text=load_fixture("ip-api.com.json"))
-
-    with patch("homeassistant.util.location._get_ipapi", return_value=None):
-        info = await location_util.async_detect_location_info(session, _test_real=True)
-
-    assert info is not None
-    assert info.ip == "1.2.3.4"
-    assert info.country_code == "US"
-    assert info.country_name == "United States"
-    assert info.region_code == "CA"
-    assert info.region_name == "California"
-    assert info.city == "San Diego"
-    assert info.zip_code == "92122"
-    assert info.time_zone == "America/Los_Angeles"
-    assert info.latitude == 32.8594
-    assert info.longitude == -117.2073
-    assert not info.use_metric
-
-
-async def test_detect_location_info_both_queries_fail(session):
-    """Ensure we return None if both queries fail."""
-    with patch("homeassistant.util.location._get_ipapi", return_value=None), patch(
-        "homeassistant.util.location._get_ip_api", return_value=None
-    ):
-        info = await location_util.async_detect_location_info(session, _test_real=True)
-    assert info is None
-
-
-async def test_freegeoip_query_raises(raising_session):
-    """Test ipapi.co query when the request to API fails."""
-    info = await location_util._get_ipapi(raising_session)
-    assert info is None
-
-
-async def test_ip_api_query_raises(raising_session):
-    """Test ip api query when the request to API fails."""
-    info = await location_util._get_ip_api(raising_session)
+async def test_whoami_query_raises(raising_session):
+    """Test whoami query when the request to API fails."""
+    info = await location_util._get_whoami(raising_session)
     assert info is None