diff --git a/app/query_normalisation.py b/app/query_normalisation.py index f5daa85..595ae21 100644 --- a/app/query_normalisation.py +++ b/app/query_normalisation.py @@ -21,10 +21,11 @@ def is_ip_address(query: str) -> bool: def resolve_domain_name(domain_name: str) -> ipaddress.IPv4Address | None: """Resolve a domain name via DNS or return None.""" try: - ip_address: ipaddress.IPv4Address = ipaddress.IPv4Address(socket.gethostbyname(domain_name)) - except socket.gaierror: - ip_address = None - return ip_address + result: str = socket.gethostbyname(domain_name) + ip_address: ipaddress.IPv4Address = ipaddress.IPv4Address(result) + return ip_address + except (socket.gaierror, ipaddress.AddressValueError): + return None def get_public_ip() -> ipaddress.IPv4Address: diff --git a/tests/test_ip_info.py b/tests/test_ip_info.py index 52af7a9..d8901e2 100644 --- a/tests/test_ip_info.py +++ b/tests/test_ip_info.py @@ -5,15 +5,15 @@ import requests_mock from app.ip_info import ( # pragma: no cover - get_ip_information, get_autonomous_system_number, + get_ip_information, get_prefix_information, ) def test_get_ip_information() -> None: """TEST: ensure that the IP information API is working correctly.""" - test_query = "1.2.3.4" + test_query: str = "1.2.3.4" ip_info = get_ip_information(test_query) assert ip_info.get("status") == "success" and ip_info.get("query") == test_query diff --git a/tests/test_query_normalisation.py b/tests/test_query_normalisation.py index 569625e..f6183ec 100644 --- a/tests/test_query_normalisation.py +++ b/tests/test_query_normalisation.py @@ -31,7 +31,7 @@ def test_resolve_domain_name_true() -> None: """TEST: Verifies that DNS resolution is working correctly.""" domain_name = "one.one.one.one" expected_results: list[str] = ["1.1.1.1", "1.0.0.1"] # Could resolve to either IP - assert resolve_domain_name(domain_name) in expected_results + assert str(resolve_domain_name(domain_name)) in expected_results def test_resolve_domain_name_false() -> None: