diff --git a/hole/exceptions.py b/hole/exceptions.py index a0c23aa..ec0025c 100644 --- a/hole/exceptions.py +++ b/hole/exceptions.py @@ -4,10 +4,19 @@ class HoleError(Exception): """General HoleError exception occurred.""" - pass + def __init__(self, message: str, *, status: int | None = None) -> None: + """Initialize the error, optionally carrying the HTTP status code.""" + super().__init__(message) + self.status = status class HoleConnectionError(HoleError): """When a connection error is encountered.""" - pass + +class HoleAuthenticationError(HoleError): + """When authentication is required, missing, or invalid (HTTP 401).""" + + +class HoleResponseError(HoleError): + """When the API returns an unexpected or unparseable payload.""" diff --git a/hole/v6.py b/hole/v6.py index ea5b1e5..c0587be 100644 --- a/hole/v6.py +++ b/hole/v6.py @@ -85,8 +85,8 @@ async def authenticate(self): ) if response.status == 401: - raise exceptions.HoleError( - "Authentication failed: Invalid password" + raise exceptions.HoleAuthenticationError( + "Authentication failed: Invalid password", status=401 ) elif response.status == 400: try: @@ -96,16 +96,21 @@ async def authenticate(self): ) except json.JSONDecodeError: error_msg = "Bad request" - raise exceptions.HoleError(f"Authentication failed: {error_msg}") + raise exceptions.HoleError( + f"Authentication failed: {error_msg}", status=400 + ) elif response.status != 200: raise exceptions.HoleError( - f"Authentication failed with status {response.status}" + f"Authentication failed with status {response.status}", + status=response.status, ) try: data = json.loads(await response.text()) except json.JSONDecodeError as err: - raise exceptions.HoleError(f"Invalid JSON response: {err}") + raise exceptions.HoleResponseError( + f"Invalid JSON response: {err}" + ) from err session_data = data.get("session", {}) @@ -195,12 +200,25 @@ async def _fetch_data(self, endpoint: str, params=None) -> dict: url, params=params, headers=headers, ssl=self.verify_tls ) + # Still unauthorized after a retry means auth is genuinely + # required and could not be satisfied with the credentials. + if response.status == 401: + raise exceptions.HoleAuthenticationError( + "Authentication required", status=401 + ) + if response.status != 200: raise exceptions.HoleError( - f"Failed to fetch data: {response.status}" + f"Failed to fetch data: {response.status}", + status=response.status, ) - return await response.json() + try: + return await response.json() + except (aiohttp.ContentTypeError, ValueError) as err: + raise exceptions.HoleResponseError( + f"Invalid response from {endpoint}" + ) from err except (asyncio.TimeoutError, aiohttp.ClientError, socket.gaierror) as err: raise exceptions.HoleConnectionError( @@ -228,7 +246,11 @@ async def get_versions(self): """Get version information from Pi-hole.""" await self.ensure_auth() response = await self._fetch_data("/info/version") - self.versions = response.get("version", {}) + if not isinstance(response, dict) or not isinstance( + response.get("version"), dict + ): + raise exceptions.HoleResponseError("Unexpected response from /info/version") + self.versions = response["version"] async def enable(self): """Enable DNS blocking.""" diff --git a/tests/test_hole.py b/tests/test_hole.py index b0d2d8c..b78e94b 100644 --- a/tests/test_hole.py +++ b/tests/test_hole.py @@ -36,6 +36,14 @@ def aiohttp_session(): return MagicMock(spec=aiohttp.ClientSession) +@pytest.fixture +def hole6(aiohttp_session): + """Fixture for a HoleV6 client with auth stubbed out.""" + client = HoleV6("localhost", aiohttp_session, password="pw") + client.ensure_auth = AsyncMock() + return client + + @pytest.mark.asyncio async def test_v5_get_data_success(aiohttp_session): """Test successful data retrieval for HoleV5.""" @@ -79,10 +87,8 @@ async def test_v5_get_data_connection_error(aiohttp_session): @pytest.mark.asyncio -async def test_v6_get_data_success(aiohttp_session): +async def test_v6_get_data_success(hole6): """Test successful data retrieval for HoleV6.""" - hole6 = HoleV6("localhost", aiohttp_session, password="pw") - hole6.ensure_auth = AsyncMock() hole6._fetch_data = AsyncMock( side_effect=[ { @@ -104,17 +110,19 @@ async def test_v6_get_data_success(aiohttp_session): {"upstreams": ["8.8.8.8"]}, {"blocking": "enabled"}, { - "core": { - "local": {"version": "1", "hash": "a"}, - "remote": {"version": "2", "hash": "b"}, - }, - "web": { - "local": {"version": "1", "hash": "a"}, - "remote": {"version": "2", "hash": "b"}, - }, - "ftl": { - "local": {"version": "1", "hash": "a"}, - "remote": {"version": "2", "hash": "b"}, + "version": { + "core": { + "local": {"version": "1", "hash": "a"}, + "remote": {"version": "2", "hash": "b"}, + }, + "web": { + "local": {"version": "1", "hash": "a"}, + "remote": {"version": "2", "hash": "b"}, + }, + "ftl": { + "local": {"version": "1", "hash": "a"}, + "remote": {"version": "2", "hash": "b"}, + }, }, }, ] @@ -131,3 +139,61 @@ async def test_v6_get_data_success(aiohttp_session): assert hole6.unique_clients == 2 assert hole6.clients_ever_seen == 3 assert hole6.status == "enabled" + + +@pytest.mark.asyncio +async def test_v6_fetch_data_authentication_error(aiohttp_session, hole6): + """Test a persistent 401 raises HoleAuthenticationError with status.""" + hole6.authenticate = AsyncMock() + aiohttp_session.get = AsyncMock(return_value=DummyResponse(status=401)) + with pytest.raises(exceptions.HoleAuthenticationError) as err: + await hole6._fetch_data("/info/version") + assert err.value.status == 401 + + +@pytest.mark.asyncio +async def test_v6_fetch_data_error_status(aiohttp_session, hole6): + """Test a non-200, non-401 status raises HoleError carrying the status.""" + aiohttp_session.get = AsyncMock(return_value=DummyResponse(status=500)) + with pytest.raises(exceptions.HoleError) as err: + await hole6._fetch_data("/stats/summary") + assert err.value.status == 500 + assert not isinstance(err.value, exceptions.HoleAuthenticationError) + + +@pytest.mark.asyncio +async def test_v6_fetch_data_invalid_response(aiohttp_session, hole6): + """Test an unparseable payload raises HoleResponseError.""" + response = DummyResponse(status=200) + response.json = AsyncMock(side_effect=ValueError("bad json")) + aiohttp_session.get = AsyncMock(return_value=response) + with pytest.raises(exceptions.HoleResponseError): + await hole6._fetch_data("/info/version") + + +@pytest.mark.asyncio +async def test_v6_get_versions_success(hole6): + """Test get_versions stores the version payload.""" + versions = {"core": {"local": {"version": "1", "hash": "a"}}} + hole6._fetch_data = AsyncMock(return_value={"version": versions}) + await hole6.get_versions() + assert hole6.versions == versions + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [[], {"version": "v6"}, {}]) +async def test_v6_get_versions_unexpected_payload(hole6, payload): + """Test an unexpected /info/version shape raises HoleResponseError.""" + hole6._fetch_data = AsyncMock(return_value=payload) + with pytest.raises(exceptions.HoleResponseError): + await hole6.get_versions() + + +@pytest.mark.asyncio +async def test_v6_authenticate_invalid_password(aiohttp_session): + """Test authenticate raises HoleAuthenticationError on 401.""" + hole6 = HoleV6("localhost", aiohttp_session, password="wrong") + aiohttp_session.post = AsyncMock(return_value=DummyResponse(status=401)) + with pytest.raises(exceptions.HoleAuthenticationError) as err: + await hole6.authenticate() + assert err.value.status == 401