Raise AuthFailedError on expired token during WebSocket reconnect - #818
Open
adam8833 wants to merge 2 commits into
Open
Raise AuthFailedError on expired token during WebSocket reconnect#818adam8833 wants to merge 2 commits into
adam8833 wants to merge 2 commits into
Conversation
WebSocketClient.init() treated a connection_error reply (sent by AppSync when the ID token has expired) the same as any other unexpected message, raising a generic WebSocketError. NiceGOApi.connect retries on WebSocketError with the same (now-expired) id_token, producing an infinite retry loop that never re-authenticates. This is the root cause of home-assistant/core#126370, where the integration silently stops working until the config entry is reloaded. Now init() inspects connection_error payloads for an UnauthorizedException and raises AuthFailedError instead, matching the handling already used for HTTP GraphQL calls. AuthFailedError isn't in connect()'s retry list, so it propagates immediately, letting the caller refresh the token instead of retrying with a dead one.
Reviewer's GuideHandle WebSocket connection_error messages indicating expired/invalid ID tokens as authentication failures, reuse a shared UnauthorizedException helper, and add tests to validate the new behavior. Sequence diagram for WebSocket initialization handling expired tokenssequenceDiagram
actor HomeAssistantCoordinator
participant NiceGOApi as NiceGOApi
participant WebSocketClient as WebSocketClient
participant AppSync as AppSync
HomeAssistantCoordinator->>NiceGOApi: connect
NiceGOApi->>WebSocketClient: init
WebSocketClient->>AppSync: WebSocket handshake
AppSync-->>WebSocketClient: connection_error
WebSocketClient->>WebSocketClient: find_unauthorized_error(errors)
alt UnauthorizedException found
WebSocketClient-->>NiceGOApi: AuthFailedError
NiceGOApi-->>HomeAssistantCoordinator: AuthFailedError
else other connection_error
WebSocketClient-->>NiceGOApi: WebSocketError
NiceGOApi->>NiceGOApi: retry on WebSocketError
end
Flow diagram for shared UnauthorizedException handlingflowchart TD
A[GraphQL or WebSocket response with errors] --> B[find_unauthorized_error]
B -->|UnauthorizedException found| C[raise AuthFailedError]
B -->|No UnauthorizedException| D{Caller}
D -->|HTTP GraphQL| E[raise ApiError]
D -->|WebSocket init non_ack| F[raise WebSocketError]
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- In
init(), you only inspect the first element oferrors; consider iterating over all returned errors to catch anUnauthorizedExceptioneven if it isn't the first. - The WebSocket error-handling logic in
init()is now similar to_check_response_errors; consider refactoring to share a common helper so the Unauthorized/other error mapping stays consistent between HTTP and WebSocket paths.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `init()`, you only inspect the first element of `errors`; consider iterating over all returned errors to catch an `UnauthorizedException` even if it isn't the first.
- The WebSocket error-handling logic in `init()` is now similar to `_check_response_errors`; consider refactoring to share a common helper so the Unauthorized/other error mapping stays consistent between HTTP and WebSocket paths.
## Individual Comments
### Comment 1
<location path="tests/test_ws_client.py" line_range="77-58" />
<code_context>
+ await mock_ws_client.init()
+
+
+async def test_ws_init_connection_error_other(mock_ws_client: WebSocketClient) -> None:
+ assert mock_ws_client.ws is not None
+ assert isinstance(mock_ws_client.ws, AsyncMock)
+ mock_ws_client.ws.receive = AsyncMock()
+ mock_ws_client.ws.receive.return_value = MagicMock(
+ data=json.dumps(
+ {
</code_context>
<issue_to_address>
**suggestion (testing):** Consider tests for other `connection_error` shapes (empty errors list, missing payload, or UnauthorizedException not in the first position)
Currently we only assert behavior when `UnauthorizedException` is the first error and when a different `errorType` is present. Please also add tests for:
- `connection_error` with an empty `errors` list
- `connection_error` with no `payload` key
- `connection_error` where `UnauthorizedException` is not the first element in `errors`
so the expected fallback behavior (e.g., using `WebSocketError` unless the first error is `UnauthorizedException`) is clearly defined and enforced.
Suggested implementation:
```python
with pytest.raises(AuthFailedError, match="Token has expired."):
await mock_ws_client.init()
async def test_ws_init_connection_error_empty_errors(
mock_ws_client: WebSocketClient,
) -> None:
assert mock_ws_client.ws is not None
assert isinstance(mock_ws_client.ws, AsyncMock)
mock_ws_client.ws.receive = AsyncMock()
mock_ws_client.ws.receive.return_value = MagicMock(
data=json.dumps(
{
"type": "connection_error",
"payload": {
"errors": [],
},
}
),
)
with pytest.raises(WebSocketError):
await mock_ws_client.init()
async def test_ws_init_connection_error_missing_payload(
mock_ws_client: WebSocketClient,
) -> None:
assert mock_ws_client.ws is not None
assert isinstance(mock_ws_client.ws, AsyncMock)
mock_ws_client.ws.receive = AsyncMock()
mock_ws_client.ws.receive.return_value = MagicMock(
data=json.dumps(
{
"type": "connection_error",
}
),
)
with pytest.raises(WebSocketError):
await mock_ws_client.init()
async def test_ws_init_connection_error_unauthorized_not_first(
mock_ws_client: WebSocketClient,
) -> None:
assert mock_ws_client.ws is not None
assert isinstance(mock_ws_client.ws, AsyncMock)
mock_ws_client.ws.receive = AsyncMock()
mock_ws_client.ws.receive.return_value = MagicMock(
data=json.dumps(
{
"type": "connection_error",
"payload": {
"errors": [
{
"errorType": "SomeOtherError",
"message": "Something else failed.",
},
{
"errorType": "UnauthorizedException",
"message": "Token has expired.",
},
],
},
}
),
)
with pytest.raises(WebSocketError):
await mock_ws_client.init()
```
1. Ensure `WebSocketError` is imported at the top of `tests/test_ws_client.py` from the same module as in the existing tests for non-`UnauthorizedException` `connection_error` handling (e.g., `from <your_package>.ws_client import WebSocketError` or equivalent).
2. If your existing tests use a different error type or message pattern for the generic case, adjust the `pytest.raises(WebSocketError)` expectations (e.g., add `match="..."`) to stay consistent with the rest of the test suite.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Sourcery pointed out init() only checked errors[0], so an UnauthorizedException later in the list would be missed. Extracted find_unauthorized_error() to _util.py and use it from both the WebSocket init() path and _check_response_errors (the existing HTTP GraphQL error path), so both scan the full errors list the same way. Non-auth connection_error messages still fall through to the existing generic WebSocketError in init(), preserving current retry behavior for cases unrelated to token expiry. Added tests for empty errors, missing payload, and an UnauthorizedException that isn't first in the list.
Author
|
Thanks for the review! Pushed a follow-up commit:
105 → 108 tests, still 100% coverage, clean ruff/mypy. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #818 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 15 15
Lines 1194 1237 +43
=========================================
+ Hits 1194 1237 +43 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Owner
|
@sourcery-ai review |
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- In
WebSocketClient.init, consider usingdata.get("type")instead ofdata["type"]to avoid aKeyErroron malformed messages and preserve the genericWebSocketErrorbehavior. - The string literal
"UnauthorizedException"is now used in multiple places; consider centralizing it as a constant to avoid typos and ease future changes if the backend error type name changes.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `WebSocketClient.init`, consider using `data.get("type")` instead of `data["type"]` to avoid a `KeyError` on malformed messages and preserve the generic `WebSocketError` behavior.
- The string literal `"UnauthorizedException"` is now used in multiple places; consider centralizing it as a constant to avoid typos and ease future changes if the backend error type name changes.
## Individual Comments
### Comment 1
<location path="tests/test_ws_client.py" line_range="77-86" />
<code_context>
+ await mock_ws_client.init()
+
+
+async def test_ws_init_connection_error_other(mock_ws_client: WebSocketClient) -> None:
+ assert mock_ws_client.ws is not None
+ assert isinstance(mock_ws_client.ws, AsyncMock)
+ mock_ws_client.ws.receive = AsyncMock()
+ mock_ws_client.ws.receive.return_value = MagicMock(
+ data=json.dumps(
+ {
+ "type": "connection_error",
+ "payload": {"errors": [{"errorType": "SomeOtherError"}]},
+ },
+ ),
+ )
+ with pytest.raises(WebSocketError):
+ await mock_ws_client.init()
+
</code_context>
<issue_to_address>
**suggestion:** Consider factoring out the repeated WebSocket mock setup into a small helper to improve test maintainability.
The tests all repeat the same setup pattern: checking `ws` is not `None`, asserting `AsyncMock`, overriding `ws.receive`, and configuring `return_value.data`. Moving this into a helper or fixture (e.g., `set_ws_receive_data(mock_ws_client, payload: dict[str, Any])`) would reduce duplication and make it easier to extend `connection_error` cases. This is an optional refactor but would improve readability and maintainability.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+77
to
+86
| async def test_ws_init_connection_error_other(mock_ws_client: WebSocketClient) -> None: | ||
| assert mock_ws_client.ws is not None | ||
| assert isinstance(mock_ws_client.ws, AsyncMock) | ||
| mock_ws_client.ws.receive = AsyncMock() | ||
| mock_ws_client.ws.receive.return_value = MagicMock( | ||
| data=json.dumps( | ||
| { | ||
| "type": "connection_error", | ||
| "payload": {"errors": [{"errorType": "SomeOtherError"}]}, | ||
| }, |
There was a problem hiding this comment.
suggestion: Consider factoring out the repeated WebSocket mock setup into a small helper to improve test maintainability.
The tests all repeat the same setup pattern: checking ws is not None, asserting AsyncMock, overriding ws.receive, and configuring return_value.data. Moving this into a helper or fixture (e.g., set_ws_receive_data(mock_ws_client, payload: dict[str, Any])) would reduce duplication and make it easier to extend connection_error cases. This is an optional refactor but would improve readability and maintainability.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Home Assistant users report the Nice G.O. integration silently stopping updates after roughly 1–7 days, only recovering after a full config entry reload or HA restart: home-assistant/core#126370
Debug logs from that issue show the WebSocket connection dying with
WebSocket connection closed, retrying, and then receiving:Root cause
WebSocketClient.init()only recognizesconnection_ack. Any other message type — includingconnection_errorsent by AppSync when the ID token has expired — falls into the generic branch and raises a plainWebSocketError.NiceGOApi.connect()'s@retrydecorator retries onWebSocketError, but nothing refreshesid_tokenin between retries, so it retries forever with the same expired token. The_check_response_errorshelper already distinguishesUnauthorizedExceptionthis way for the HTTP GraphQL calls (get_all_barriers,open_barrier, etc.), but the WebSocket handshake path never got the same treatment.Fix
init()now inspectsconnection_errorpayloads forUnauthorizedExceptionand raisesAuthFailedErrorinstead ofWebSocketError.AuthFailedErrorisn't inconnect()'s retry list, so it propagates immediately instead of being silently retried — letting the caller (e.g. Home Assistant's coordinator) refresh the token and reconnect, rather than hammering AppSync with a dead token indefinitely.Testing
test_ws_init_token_expiredandtest_ws_init_connection_error_othertotests/test_ws_client.py.poetry run pytest— 105 passed, 100% coverage.poetry run ruff check/ruff format --check/mypy— clean.Follow-up
Home Assistant's
nice_gocoordinator (client_listen) would need a small change once this ships, to catchAuthFailedErrorthere and callauthenticate_refreshbefore retryingconnect()— happy to help with that PR once this is released.Summary by Sourcery
Handle expired or invalid ID tokens during WebSocket initialization by surfacing an authentication failure instead of endlessly retrying.
Bug Fixes:
Enhancements:
Tests: