diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..9b388533 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,7 @@ +{ + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true +} \ No newline at end of file diff --git a/README.md b/README.md index eec5e386..e2e1dd68 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ > However, it is not going to be deprecated according to [this comment](https://github.com/home-assistant/developers.home-assistant/pull/2150#pullrequestreview-2017433583) > But it is recommended to use the Websocket API for new integrations. -Here is a quick example. +### REST API Examples ```py from homeassistant_api import Client @@ -25,7 +25,7 @@ with Client( '', # i.e. 'http://homeassistant.local:8123/api/' '' ) as client: - light = client.trigger_service('light', 'turn_on', {'entity_id': 'light.living_room'}) + light = client.trigger_service('light', 'turn_on', entity_id="light.living_room") ``` All the methods also support async/await! @@ -33,6 +33,35 @@ Just prefix the method with `async_` and pass the `use_async=True` argument to t Then you can use the methods as coroutines (i.e. `await light.async_turn_on(...)`). +```py +import asyncio +from homeassistant_api import Client + +async def main(): + with Client( + '', # i.e. 'http://homeassistant.local:8123/api/' + '', + use_async=True + ) as client: + light = await client.async_trigger_service('light', 'turn_on', entity_id="light.living_room") + +asyncio.run(main()) +``` + +### Websocket API Example + +```py +from homeassistant_api import WebsocketClient + +with WebsocketClient( + '', # i.e. 'ws://homeassistant.local:8123/api/websocket' + '' +) as ws_client: + light = ws_client.trigger_service('light', 'turn_on', entity_id="light.living_room") +``` + +> Note: The Websocket API is not yet supported in async/await mode. + ## Documentation All documentation, API reference, contribution guidelines and pretty much everything else diff --git a/compose.yml b/compose.yml index 405c84dd..2a7d627c 100644 --- a/compose.yml +++ b/compose.yml @@ -1,5 +1,3 @@ -version: "3.9" - services: server: image: "homeassistant/home-assistant:stable" @@ -19,6 +17,7 @@ services: - server environment: HOMEASSISTANTAPI_URL: http://server:8123/api + HOMEASSISTANTAPI_WS_URL: ws://server:8123/api/websocket HOMEASSISTANTAPI_TOKEN: ${HOMEASSISTANTAPI_TOKEN} DELAY: 60 diff --git a/docs/api.rst b/docs/api.rst index 36ba9ac7..47ab8172 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -4,4 +4,4 @@ Code Reference .. automodule:: homeassistant_api :platform: Linux, Windows, MacOS :inherited-members: - :exclude-members: model_json_schema, model_copy, model_rebuild, model_dump, construct, copy, dict, from_orm, json, parse_file, parse_obj, parse_raw, parse_str, parse_url, schema, schema_json, schema_yaml, schema_yml, to_orm, update_forward_refs, validate, validate_file, validate_obj, validate_raw, validate_str, validate_url, model_validate_strings, model_validate_json, model_validate, model_post_init, model_parametrized_name, model_extra, model_fields_set, model_dump_json, model_construct, model_computed_fields \ No newline at end of file + :exclude-members: model_json_schema, model_copy, model_rebuild, model_dump, construct, copy, dict, from_orm, json, parse_file, model_validate, parse_raw, parse_str, parse_url, schema, schema_json, schema_yaml, schema_yml, to_orm, update_forward_refs, validate, validate_file, validate_obj, validate_raw, validate_str, validate_url, model_validate_strings, model_validate_json, model_validate, model_post_init, model_parametrized_name, model_extra, model_fields_set, model_dump_json, model_construct, model_computed_fields \ No newline at end of file diff --git a/homeassistant_api/__init__.py b/homeassistant_api/__init__.py index 8f6ca386..da79af0a 100644 --- a/homeassistant_api/__init__.py +++ b/homeassistant_api/__init__.py @@ -20,6 +20,7 @@ "ParameterMissingError", "RequestError", "UnauthorizedError", + "WebsocketClient", ) from .client import Client @@ -36,6 +37,7 @@ ) from .models import Domain, Entity, Event, Group, History, LogbookEntry, Service, State from .processing import Processing +from .websocket import WebsocketClient Domain.model_rebuild() Entity.model_rebuild() diff --git a/homeassistant_api/client.py b/homeassistant_api/client.py index 94c23c86..c5ddd6eb 100644 --- a/homeassistant_api/client.py +++ b/homeassistant_api/client.py @@ -1,9 +1,8 @@ """Module containing the primary Client class.""" import logging -from typing import Any import urllib.parse as urlparse -import warnings +from typing import Any from .rawasyncclient import RawAsyncClient from .rawclient import RawClient @@ -41,9 +40,5 @@ def __init__( RawClient.__init__( self, api_url, token, verify_ssl=verify_ssl, **kwargs ) - warnings.warn( - "The REST API is being phased out and will be removed in a far future release. Please use the WebSocket API instead.", - DeprecationWarning, - ) else: raise ValueError(f"Unknown scheme {parsed.scheme} in {api_url}") diff --git a/homeassistant_api/errors.py b/homeassistant_api/errors.py index 015c93ad..f3fe493f 100644 --- a/homeassistant_api/errors.py +++ b/homeassistant_api/errors.py @@ -1,6 +1,6 @@ """Module for custom error classes""" -from typing import Union +from typing import Optional, Union class HomeassistantAPIError(Exception): @@ -55,8 +55,8 @@ def __init__(self, status_code: int, content: Union[str, bytes]) -> None: class UnauthorizedError(HomeassistantAPIError): """Error raised when an invalid token in used to authenticate with homeassistant.""" - def __init__(self) -> None: - super().__init__("Invalid authentication token") + def __init__(self, message: Optional[str] = None) -> None: + super().__init__(message or "Invalid authentication token") class EndpointNotFoundError(HomeassistantAPIError): diff --git a/homeassistant_api/models/base.py b/homeassistant_api/models/base.py index 92b9b418..e003683d 100644 --- a/homeassistant_api/models/base.py +++ b/homeassistant_api/models/base.py @@ -3,8 +3,8 @@ from datetime import datetime from typing import Annotated -from pydantic import ConfigDict, BaseModel as PydanticBaseModel, PlainSerializer - +from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict, PlainSerializer DatetimeIsoField = Annotated[ datetime, diff --git a/homeassistant_api/models/domains.py b/homeassistant_api/models/domains.py index 5af69f17..82d68441 100644 --- a/homeassistant_api/models/domains.py +++ b/homeassistant_api/models/domains.py @@ -12,19 +12,24 @@ from .states import State if TYPE_CHECKING: - from homeassistant_api import Client + from homeassistant_api import Client, WebsocketClient class Domain(BaseModel): """Model representing the domain that services belong to.""" - def __init__(self, *args, _client: Optional["Client"] = None, **kwargs) -> None: + def __init__( + self, + *args, + _client: Optional[Union["Client", "WebsocketClient"]] = None, + **kwargs, + ) -> None: super().__init__(*args, **kwargs) if _client is None: raise ValueError("No client passed.") object.__setattr__(self, "_client", _client) - _client: "Client" + _client: Union["Client", "WebsocketClient"] domain_id: str = Field( ..., description="The name of the domain that services belong to. " @@ -36,7 +41,9 @@ def __init__(self, *args, _client: Optional["Client"] = None, **kwargs) -> None: ) @classmethod - def from_json(cls, json: Dict[str, Any], client: "Client") -> "Domain": + def from_json( + cls, json: Dict[str, Any], client: Union["Client", "WebsocketClient"] + ) -> "Domain": """Constructs Domain and Service models from json data.""" if "domain" not in json or "services" not in json: raise ValueError("Missing services or domain attribute in json argument.") @@ -96,10 +103,15 @@ class Service(BaseModel): description: Optional[str] = None fields: Optional[Dict[str, ServiceField]] = None - def trigger( - self, **service_data - ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]: + def trigger(self, entity_id: Optional[str] = None, **service_data) -> Union[ + Tuple[State, ...], + Tuple[Tuple[State, ...], Dict[str, Any]], + dict[str, Any], + None, + ]: """Triggers the service associated with this object.""" + if entity_id is not None: + service_data["entity_id"] = entity_id try: return self.domain._client.trigger_service_with_response( self.domain.domain_id, @@ -114,9 +126,18 @@ def trigger( ) async def async_trigger( - self, **service_data + self, entity_id: Optional[str] = None, **service_data ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]: """Triggers the service associated with this object.""" + if entity_id is not None: + service_data["entity_id"] = entity_id + + from homeassistant_api import WebsocketClient # prevent circular import + + if isinstance(self.domain._client, WebsocketClient): + raise NotImplementedError( + "WebsocketClient does not support async/await syntax." + ) try: return await self.domain._client.async_trigger_service_with_response( self.domain.domain_id, @@ -130,8 +151,13 @@ async def async_trigger( **service_data, ) - def __call__(self, **service_data) -> Union[ - Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]], + def __call__(self, entity_id: Optional[str] = None, **service_data) -> Union[ + Union[ + Tuple[State, ...], + Tuple[Tuple[State, ...], Dict[str, Any]], + dict[str, Any], + None, + ], Coroutine[ Any, Any, Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]] ], @@ -145,7 +171,7 @@ def __call__(self, **service_data) -> Union[ if inspect.iscoroutinefunction( caller := gc.get_referrers(parent_frame.f_code)[0] ) or inspect.iscoroutine(caller): - return self.async_trigger(**service_data) + return self.async_trigger(entity_id=entity_id, **service_data) except IndexError: # pragma: no cover pass - return self.trigger(**service_data) + return self.trigger(entity_id=entity_id, **service_data) diff --git a/homeassistant_api/models/states.py b/homeassistant_api/models/states.py index 0ff01844..e5f8c272 100644 --- a/homeassistant_api/models/states.py +++ b/homeassistant_api/models/states.py @@ -12,9 +12,22 @@ class Context(BaseModel): """Model for entity state contexts.""" id: str = Field( - max_length=128, + max_length=128, # arbitrary limit description="Unique string identifying the context.", ) + parent_id: Optional[str] = Field( + max_length=128, + description="Unique string identifying the parent context.", + ) + user_id: Optional[str] = Field( + max_length=128, + description="Unique string identifying the user.", + ) + + @classmethod + def from_json(cls, json: Dict[str, Any]) -> "Context": + """Constructs Context model from json data""" + return cls.model_validate(json) class State(BaseModel): diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py new file mode 100644 index 00000000..7aae4b38 --- /dev/null +++ b/homeassistant_api/models/websocket.py @@ -0,0 +1,97 @@ +"""A module defining the responses we expect from the websocket API.""" + +from typing import Any, Literal, Optional, Union + +from .base import BaseModel +from .states import Context, DatetimeIsoField + +__all__ = ( + "AuthRequired", + "AuthOk", + "AuthInvalid", + "PingResponse", + "ErrorResponse", + "ResultResponse", + "EventResponse", +) + + +class AuthRequired(BaseModel): + type: Literal["auth_required"] + ha_version: str + + +class AuthOk(BaseModel): + type: Literal["auth_ok"] + ha_version: str + + +class AuthInvalid(BaseModel): + type: Literal["auth_invalid"] + message: str + + +class PingResponse(BaseModel): + """Ping websocket response model.""" + + id: int + type: Literal["pong"] + start: int # added by the client, nanoseconds + end: Optional[int] = None # added by the client, nanoseconds + + +class Error(BaseModel): + code: str + message: str + + +class ErrorResponse(BaseModel): + """Error websocket response model.""" + + id: int + success: Literal[False] + type: Literal["result"] + error: Error + + +class ResultResponse(BaseModel): + """Result websocket response model.""" + + id: int + success: Literal[True] + type: Literal["result"] + result: Optional[Any] + + +class FiredEvent(BaseModel): + """A model to parse the `event` key of fired event websocket responses.""" + + event_type: str + data: dict[str, Any] + + origin: Literal["LOCAL", "REMOTE"] + # REMOTE if another API client or webhook fired the event + # LOCAL if Home Assistant (or the auth token we used) fired the event + + time_fired: DatetimeIsoField # datetime.datetime + context: Optional[Context] + + +class TemplateEvent(BaseModel): + result: str + listeners: dict[str, Any] + + +class FiredTrigger(BaseModel): + """A model to parse the `trigger` key of fired event websocket responses.""" + + context: Optional[Context] + variables: dict[str, Any] + + +class EventResponse(BaseModel): + """A model to parse the response of a fired event websocket response.""" + + id: int + type: Literal["event"] + event: Union[FiredEvent, FiredTrigger, TemplateEvent] diff --git a/homeassistant_api/rawasyncclient.py b/homeassistant_api/rawasyncclient.py index 18c6197c..778e753b 100644 --- a/homeassistant_api/rawasyncclient.py +++ b/homeassistant_api/rawasyncclient.py @@ -27,6 +27,7 @@ from .models import Domain, Entity, Event, Group, History, LogbookEntry, State from .processing import AsyncResponseType, Processing from .rawbaseclient import RawBaseClient +from .utils import prepare_entity_id if TYPE_CHECKING: from homeassistant_api import Client @@ -144,7 +145,7 @@ async def async_get_logbook_entries( params, url = self.prepare_get_logbook_entry_params(*args, **kwargs) data = await self.async_request(url, params=params) for entry in data: - yield LogbookEntry.parse_obj(entry) + yield LogbookEntry.model_validate(entry) async def async_get_entity_histories( self, @@ -169,7 +170,7 @@ async def async_get_entity_histories( params=self.construct_params(params), ) for states in data: - yield History.parse_obj({"states": states}) + yield History.model_validate({"states": states}) async def async_get_rendered_template(self, template: str) -> str: """ @@ -232,9 +233,9 @@ async def async_get_entities(self) -> Dict[str, Group]: async def async_get_entity( self, - group_id: str | None = None, - slug: str | None = None, - entity_id: str | None = None, + group_id: Optional[str] = None, + slug: Optional[str] = None, + entity_id: Optional[str] = None, ) -> Optional[Entity]: """ Returns a Entity model for an :code:`entity_id`. @@ -335,7 +336,7 @@ async def async_get_state( # pylint: disable=duplicate-code Fetches the state of the entity specified. :code:`GET /api/states/` """ - target_entity_id = self.prepare_entity_id( + target_entity_id = prepare_entity_id( group_id=group_id, slug=slug, entity_id=entity_id, @@ -355,7 +356,7 @@ async def async_set_state( # pylint: disable=duplicate-code data = await self.async_request( join("states", state.entity_id), method="POST", - json=json.loads(state.json()), + json=json.loads(state.model_dump_json()), ) return State.from_json(cast(Dict[Any, Any], data)) diff --git a/homeassistant_api/rawbaseclient.py b/homeassistant_api/rawbaseclient.py index b5c2a8af..6e506983 100644 --- a/homeassistant_api/rawbaseclient.py +++ b/homeassistant_api/rawbaseclient.py @@ -1,9 +1,8 @@ """Module for parent RawWrapper class""" -import re from datetime import datetime from posixpath import join -from typing import Dict, Iterable, Optional, Tuple, Union, Any +from typing import Any, Dict, Iterable, Optional, Tuple, Union from .models import Entity @@ -66,36 +65,6 @@ def construct_params(params: Dict[str, Optional[str]]) -> str: """Custom method for constructing non-standard query strings""" return "&".join([k if v is None else f"{k}={v}" for k, v in params.items()]) - @staticmethod - def format_entity_id(entity_id: str) -> str: - """Takes in a string and formats it into valid snake_case.""" - entity_id = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", entity_id) - entity_id = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", entity_id) - entity_id = entity_id.replace("-", "_") - return entity_id.lower() - - def prepare_entity_id( - self, - *, - group_id: Optional[str] = None, - slug: Optional[str] = None, - entity_id: Optional[str] = None, - ) -> str: - """ - Combines optional :code:`group` and :code:`slug` into an :code:`entity_id` if provided. - Favors :code:`entity_id` over :code:`group` or :code:`slug`. - """ - if (group_id is None or slug is None) and entity_id is None: - raise ValueError( - "To use group or slug you need to pass both, not just one. " - "Otherwise pass entity_id. " - "Also make sure you are using keyword arguments." - ) - if group_id is not None and slug is not None: - entity_id = f"{group_id}.{slug}" - assert entity_id is not None - return self.format_entity_id(entity_id) - @staticmethod def prepare_get_entity_histories_params( entities: Optional[Tuple[Entity, ...]] = None, diff --git a/homeassistant_api/rawclient.py b/homeassistant_api/rawclient.py index 2f4b92e8..053bb65d 100644 --- a/homeassistant_api/rawclient.py +++ b/homeassistant_api/rawclient.py @@ -22,6 +22,8 @@ import requests import requests_cache +from homeassistant_api.utils import prepare_entity_id + from .errors import BadTemplateError, RequestError, RequestTimeoutError from .models import Domain, Entity, Event, Group, History, LogbookEntry, State from .processing import Processing, ResponseType @@ -85,7 +87,7 @@ def request( self, path: str, method="GET", - headers: Dict[str, str] | None = None, + headers: Optional[Dict[str, str]] = None, decode_bytes: bool = True, **kwargs, ) -> Any: @@ -139,7 +141,7 @@ def get_logbook_entries( params, url = self.prepare_get_logbook_entry_params(*args, **kwargs) data = self.request(url, params=params) for entry in data: - yield LogbookEntry.parse_obj(entry) + yield LogbookEntry.model_validate(entry) def get_entity_histories( self, @@ -164,7 +166,7 @@ def get_entity_histories( params=self.construct_params(params), ) for states in data: - yield History.parse_obj({"states": states}) + yield History.model_validate({"states": states}) def get_rendered_template(self, template: str) -> str: """ @@ -226,9 +228,9 @@ def get_entities(self) -> Dict[str, Group]: def get_entity( self, - group_id: str | None = None, - slug: str | None = None, - entity_id: str | None = None, + group_id: Optional[str] = None, + slug: Optional[str] = None, + entity_id: Optional[str] = None, ) -> Optional[Entity]: """ Returns an :py:class:`Entity` model for an :code:`entity_id`. @@ -331,7 +333,7 @@ def get_state( # pylint: disable=duplicate-code Fetches the state of the entity specified. :code:`GET /api/states/` """ - entity_id = self.prepare_entity_id( + entity_id = prepare_entity_id( group_id=group_id, slug=slug, entity_id=entity_id, @@ -351,7 +353,7 @@ def set_state( # pylint: disable=duplicate-code data = self.request( join("states", state.entity_id), method="POST", - json=json.loads(state.json()), + json=json.loads(state.model_dump_json()), ) return State.from_json(cast(Dict[str, Any], data)) diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py new file mode 100644 index 00000000..4f6e26ff --- /dev/null +++ b/homeassistant_api/rawwebsocket.py @@ -0,0 +1,200 @@ +import json +import logging +import time +from typing import Any, Optional, Union, cast + +import websockets.sync.client as ws +from pydantic import ValidationError + +from homeassistant_api.errors import ( + ReceivingError, + RequestError, + ResponseError, + UnauthorizedError, +) +from homeassistant_api.models.websocket import ( + AuthInvalid, + AuthOk, + AuthRequired, + ErrorResponse, + EventResponse, + PingResponse, + ResultResponse, +) + + +logger = logging.getLogger(__name__) + + +class RawWebsocketClient: + api_url: str + token: str + _conn: Optional[ws.ClientConnection] + + def __init__( + self, + api_url: str, + token: str, + ) -> None: + self.api_url = api_url + self.token = token + self._conn = None + + self._id_counter = 0 + self._result_responses: dict[int, Optional[ResultResponse]] = ( + {} + ) # id -> response + self._event_responses: dict[int, list[EventResponse]] = ( + {} + ) # id -> [response, ...] + self._ping_responses: dict[int, PingResponse] = {} # id -> (sent, received) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.api_url!r})" + + def __enter__(self): + self._conn = ws.connect(self.api_url) + self._conn.__enter__() + okay = self.authentication_phase() + logging.info("Authenticated with Home Assistant (%s)", okay.ha_version) + self.supported_features_phase() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self._conn.__exit__(exc_type, exc_value, traceback) + self._conn = None + + def _request_id(self) -> int: + """Get a unique id for a message.""" + self._id_counter += 1 + return self._id_counter + + def _send(self, data: dict[str, Any]) -> None: + """Send a message to the websocket server.""" + logger.debug(f"Sending message: {data}") + if self._conn is None: + raise ReceivingError("Connection is not open!") + self._conn.send(json.dumps(data)) + + def _recv(self) -> dict[str, Any]: + """Receive a message from the websocket server.""" + if self._conn is None: + raise ReceivingError("Connection is not open!") + _bytes = self._conn.recv() + logger.debug("Received message: %s", _bytes) + return json.loads(_bytes) + + def send(self, type: str, include_id: bool = True, **data: Any) -> int: + """ + Send a command message to the websocket server and wait for a "result" response. + + Returns the id of the message sent. + """ + if include_id: # auth messages don't have an id + data["id"] = self._request_id() + data["type"] = type + + self._send(data) + + if "id" in data: + if data["type"] == "ping": + self._ping_responses[data["id"]] = PingResponse( + start=time.perf_counter_ns(), + id=data["id"], + type="pong", + ) + else: + self._event_responses[data["id"]] = [] + self._result_responses[data["id"]] = None + return data["id"] + return -1 # non-command messages don't have an id + + def check_success(self, data: dict[str, Any]) -> None: + """Check if a command message was successful.""" + try: + error_resp = ErrorResponse.model_validate(data) + raise RequestError(error_resp.error.code, error_resp.error.message) + except ValidationError: + pass + + def handle_recv(self, data: dict[str, Any]) -> None: + """Handle a received message.""" + if "id" not in data: + raise ReceivingError( + "Received a message without an id outside the auth phase." + ) + self.check_success(data) + self.parse_response(data) + + def parse_response(self, data: dict[str, Any]) -> None: + if data.get("type") == "pong": + logger.info("Received pong message") + self._ping_responses[data["id"]].end = time.perf_counter_ns() + elif data.get("type") == "result": + logger.info("Received result message") + self._result_responses[data["id"]] = ResultResponse.model_validate(data) + elif data.get("type") == "event": + logger.info("Received event message %s", data["event"]) + self._event_responses[data["id"]].append(EventResponse.model_validate(data)) + else: + raise ReceivingError(f"Received unexpected message type: {data}") + + def recv(self, id: int) -> Union[EventResponse, ResultResponse, PingResponse]: + """Receive a response to a message from the websocket server.""" + while True: + ## have we received a message with the id we're looking for? + if self._result_responses.get(id) is not None: + return cast(dict[int, ResultResponse], self._result_responses).pop( + id + ) # ughhh why can't mypy figure this out + if self._event_responses.get(id, []): + return self._event_responses[id].pop(0) + if self._ping_responses.get(id) is not None: + if self._ping_responses[id].end is not None: + return self._ping_responses.pop(id) + + ## if not, keep receiving messages until we do + self.handle_recv(self._recv()) + + def authentication_phase(self) -> AuthOk: + """Authenticate with the websocket server.""" + # Capture the first message from the server saying we need to authenticate + try: + welcome = AuthRequired.model_validate(self._recv()) + logger.debug(f"Received welcome message: {welcome}") + except ValidationError as e: + raise ResponseError("Unexpected response during authentication") from e + + # Send our authentication token + self.send("auth", access_token=self.token, include_id=False) + logger.debug("Sent auth message") + + # Check the response + resp = self._recv() + try: + return AuthOk.model_validate(resp) + except ValidationError as e: + error_resp = AuthInvalid.model_validate(resp) + raise UnauthorizedError(error_resp.message) from e + except Exception as e: + raise ResponseError( + "Unexpected response during authentication", resp["message"] + ) from e + + def supported_features_phase(self) -> None: + """Get the supported features from the websocket server.""" + resp = self.recv( + self.send( + "supported_features", + features={ + # "coalesce_messages": 42, # including this key sets it to True + }, + ) + ) + assert cast(ResultResponse, resp).result is None + + def ping_latency(self) -> float: + """Get the latency (in milliseconds) of the connection by sending a ping message.""" + pong = cast(PingResponse, self.recv(self.send("ping"))) + assert pong.end is not None + return (pong.end - pong.start) / 1_000_000 diff --git a/homeassistant_api/utils.py b/homeassistant_api/utils.py new file mode 100644 index 00000000..ff2947ce --- /dev/null +++ b/homeassistant_api/utils.py @@ -0,0 +1,32 @@ +import re +from typing import Optional + + +def format_entity_id(entity_id: str) -> str: + """Takes in a string and formats it into valid snake_case.""" + entity_id = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", entity_id) + entity_id = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", entity_id) + entity_id = entity_id.replace("-", "_") + return entity_id.lower() + + +def prepare_entity_id( + *, + group_id: Optional[str] = None, + slug: Optional[str] = None, + entity_id: Optional[str] = None, +) -> str: + """ + Combines optional :code:`group` and :code:`slug` into an :code:`entity_id` if provided. + Favors :code:`entity_id` over :code:`group` or :code:`slug`. + """ + if (group_id is None or slug is None) and entity_id is None: + raise ValueError( + "To use group or slug you need to pass both, not just one. " + "Otherwise pass entity_id. " + "Also make sure you are using keyword arguments." + ) + if group_id is not None and slug is not None: + entity_id = f"{group_id}.{slug}" + assert entity_id is not None + return format_entity_id(entity_id) diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py new file mode 100644 index 00000000..8068ade2 --- /dev/null +++ b/homeassistant_api/websocket.py @@ -0,0 +1,293 @@ +import contextlib +import logging +import urllib.parse as urlparse +from typing import Any, Dict, Generator, Optional, Tuple, Union, cast + +from homeassistant_api.models import Domain, Entity, Group, State +from homeassistant_api.models.states import Context +from homeassistant_api.models.websocket import ( + EventResponse, + FiredEvent, + FiredTrigger, + ResultResponse, + TemplateEvent, +) +from homeassistant_api.utils import prepare_entity_id + +from .rawwebsocket import RawWebsocketClient + +logger = logging.getLogger(__name__) + + +class WebsocketClient(RawWebsocketClient): + def __init__( + self, + api_url: str, + token: str, + ) -> None: + parsed = urlparse.urlparse(api_url) + + if parsed.scheme not in {"ws", "wss"}: + raise ValueError(f"Unknown scheme {parsed.scheme} in {api_url}") + super().__init__(api_url, token) + logger.debug(f"WebSocketClient initialized with api_url: {api_url}") + + def get_rendered_template(self, template: str) -> str: + """ + Renders a Jinja2 template with Home Assistant context data. + See https://www.home-assistant.io/docs/configuration/templating. + :code:`"type": "render_template"` + """ + id = self.send("render_template", template=template, report_errors=True) + first = self.recv(id) + assert cast(ResultResponse, first).result is None + second = self.recv(id) + self._unsubscribe(id) + return cast(TemplateEvent, cast(EventResponse, second).event).result + + def get_config(self) -> dict[str, Any]: + """Get the Home Assistant configuration.""" + return cast( + dict[str, Any], + cast( + ResultResponse, + self.recv(self.send("get_config")), + ).result, + ) + + def get_states(self) -> Tuple[State, ...]: + """Get a list of states.""" + return tuple( + State.from_json(state) + for state in cast( + list[dict[str, Any]], + cast(ResultResponse, self.recv(self.send("get_states"))).result, + ) + ) + + def get_state( # pylint: disable=duplicate-code + self, + *, + entity_id: Optional[str] = None, + group_id: Optional[str] = None, + slug: Optional[str] = None, + ) -> State: + """ + Just calls the `get_states` method and filters the result. + + Please tell home-assistant/core to add a `get_state` command to the WS API! + """ + entity_id = prepare_entity_id( + group_id=group_id, + slug=slug, + entity_id=entity_id, + ) + + for state in self.get_states(): + if state.entity_id == entity_id: + return state + raise ValueError(f"Entity {entity_id} not found!") + + def get_entities(self) -> Dict[str, Group]: + """ + Fetches all entities from the Websocket API and returns them as a dictionary of :py:class:`Group`'s. + """ + entities: Dict[str, Group] = {} + for state in self.get_states(): + group_id, entity_slug = state.entity_id.split(".") + if group_id not in entities: + entities[group_id] = Group( + group_id=group_id, + _client=self, # type: ignore[arg-type] + ) + entities[group_id]._add_entity(entity_slug, state) + return entities + + def get_entity( + self, + group_id: Optional[str] = None, + slug: Optional[str] = None, + entity_id: Optional[str] = None, + ) -> Optional[Entity]: + """ + Returns an :py:class:`Entity` model for an :code:`entity_id`. + + Calls :py:meth:`get_state` in the process. + + Please tell home-assistant/core to add a `get_state` command to the WS API! + """ + if group_id is not None and slug is not None: + state = self.get_state(group_id=group_id, slug=slug) + elif entity_id is not None: + state = self.get_state(entity_id=entity_id) + else: + help_msg = ( + "Use keyword arguments to pass entity_id. " + "Or you can pass the group_id and slug instead" + ) + raise ValueError( + f"Neither group_id and slug or entity_id provided. {help_msg}" + ) + split_group_id, split_slug = state.entity_id.split(".") + group = Group( + group_id=split_group_id, + _client=self, # type: ignore[arg-type] + ) + group._add_entity(split_slug, state) + return group.get_entity(split_slug) + + def get_domains(self) -> dict[str, Domain]: + """Get a list of (service) domains.""" + resp = self.recv(self.send("get_services")) + domains = map( + lambda item: Domain.from_json( + {"domain": item[0], "services": item[1]}, + client=self, + ), + cast(dict[str, Any], cast(ResultResponse, resp).result).items(), + ) + return {domain.domain_id: domain for domain in domains} + + def get_domain(self, domain: str) -> Domain: + """Get a domain. + + Note: This is not a method in the WS API client... yet. + + Please tell home-assistant/core to add a `get_domain` command to the WS API! + + For now, just call the `get_services` method and parsing the result. + """ + return self.get_domains()[domain] + + def trigger_service( + self, + domain: str, + service: str, + entity_id: Optional[str] = None, + **service_data, + ) -> None: + """Trigger a service.""" + params = { + "domain": domain, + "service": service, + "service_data": service_data, + "return_response": False, + } + if entity_id is not None: + params["target"] = {"entity_id": entity_id} + + data = self.recv(self.send("call_service", include_id=True, **params)) + + # TODO: handle data["result"]["context"] ? + + assert ( + cast( + dict[str, Any], + cast(ResultResponse, data).result, + ).get("response") + is None + ) # should always be None for services without a response + + def trigger_service_with_response( + self, + domain: str, + service: str, + entity_id: Optional[str] = None, + **service_data, + ) -> dict[str, Any]: + params = { + "domain": domain, + "service": service, + "service_data": service_data, + "return_response": True, + } + if entity_id is not None: + params["target"] = {"entity_id": entity_id} + + data = self.recv(self.send("call_service", include_id=True, **params)) + + return cast(dict[str, Any], cast(ResultResponse, data).result)["response"] + + @contextlib.contextmanager + def subscribe_events( + self, + event_type: Optional[str] = None, + ) -> Generator[Generator[FiredEvent, None, None], None, None]: + """ + Subscribe to all events of a certain type and calls `unsubscribe_events` when done. + """ + subscription = self._subscribe_events(event_type) + yield cast(Generator[FiredEvent, None, None], self._wait_for(subscription)) + self._unsubscribe(subscription) + + def _subscribe_events(self, event_type: Optional[str]) -> int: + """Subscribe to all events of a certain type.""" + params = {"event_type": event_type} if event_type else {} + return self.recv(self.send("subscribe_events", include_id=True, **params)).id + + @contextlib.contextmanager + def subscribe_trigger( + self, trigger: str, **trigger_fields + ) -> Generator[Generator[dict[str, Any], None, None], None, None]: + """ + Subscribe to a Home Assistant trigger. + Allows additional trigger keyword parameters with **kwargs (i.e. passing `tag_id=...` for NFC tag triggers). + + Ex. + ``` + - trigger: state + entity_id: light.kitchen + ``` -> `subscribe_trigger("state", entity_id="light.kitchen")` + """ + subscription = self._subscribe_trigger(trigger, **trigger_fields) + yield ( + fired_trigger.variables + for fired_trigger in cast( + Generator[FiredTrigger, None, None], + self._wait_for(subscription), + ) + ) + self._unsubscribe(subscription) + + def _subscribe_trigger(self, trigger: str, **trigger_fields) -> int: + """Return the subscription id of the trigger we subscribe to.""" + return self.recv( + self.send( + "subscribe_trigger", trigger={"platform": trigger, **trigger_fields} + ) + ).id + + def _wait_for( + self, subscription_id: int + ) -> Generator[Union[FiredEvent, FiredTrigger], None, None]: + """ + An iterator that waits for events of a certain type. + """ + while True: + yield cast( + Union[ + FiredEvent, FiredTrigger + ], # we can cast this because TemplateEvent is only used for rendering templates + cast(EventResponse, self.recv(subscription_id)).event, + ) + + def _unsubscribe(self, subcription_id: int) -> None: + """Unsubscribe from all events of a certain type.""" + resp = self.recv(self.send("unsubscribe_events", subscription=subcription_id)) + assert cast(ResultResponse, resp).result is None + self._event_responses.pop(subcription_id) + + def fire_event(self, event_type: str, **event_data) -> Context: + """Fire an event.""" + params: dict[str, Any] = {"event_type": event_type} + if event_data: + params["event_data"] = event_data + return Context.from_json( + cast( + dict[str, dict[str, Any]], + cast( + ResultResponse, + self.recv(self.send("fire_event", include_id=True, **params)), + ).result, + )["context"] + ) diff --git a/tests/conftest.py b/tests/conftest.py index efcdc545..c9ea1995 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,12 +1,12 @@ import asyncio import logging import os -from typing import AsyncGenerator, Generator +from typing import AsyncGenerator, Generator, Literal import pytest import pytest_asyncio -from homeassistant_api import Client +from homeassistant_api import Client, WebsocketClient TIMEOUT = 300 @@ -42,7 +42,9 @@ def event_loop(): @pytest_asyncio.fixture(name="async_cached_client", scope="session") -async def setup_async_cached_client(wait_for_server) -> AsyncGenerator[Client, None]: +async def setup_async_cached_client( + wait_for_server: Literal[None], +) -> AsyncGenerator[Client, None]: """Initializes the Client and enters an async cached session.""" async with Client( os.environ["HOMEASSISTANTAPI_URL"], @@ -50,3 +52,15 @@ async def setup_async_cached_client(wait_for_server) -> AsyncGenerator[Client, N use_async=True, ) as client: yield client + + +@pytest.fixture(name="websocket_client", scope="session") +def setup_websocket_client( + wait_for_server: Literal[None], +) -> Generator[Client, None, None]: + """Initializes the Client and enters a WebSocket session.""" + with WebsocketClient( + os.environ["HOMEASSISTANTAPI_WS_URL"], + os.environ["HOMEASSISTANTAPI_TOKEN"], + ) as client: + yield client diff --git a/tests/test_client.py b/tests/test_client.py index 6c505b55..105db369 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -3,7 +3,7 @@ import aiohttp_client_cache import requests_cache -from homeassistant_api import Client +from homeassistant_api import Client, WebsocketClient def test_custom_cached_session() -> None: @@ -47,3 +47,11 @@ async def test_default_async_session() -> None: use_async=True, ): pass + + +def test_websocket_client_ping() -> None: + with WebsocketClient( + os.environ["HOMEASSISTANTAPI_WS_URL"], + os.environ["HOMEASSISTANTAPI_TOKEN"], + ) as client: + assert client.ping_latency() > 0 diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py index 587fc140..cbc1b01b 100644 --- a/tests/test_endpoints.py +++ b/tests/test_endpoints.py @@ -1,10 +1,12 @@ """Module for making sure endpoints that should succeed, do indeed succeed.""" +import logging from datetime import datetime from homeassistant_api import Client from homeassistant_api.models.events import Event from homeassistant_api.models.states import State +from homeassistant_api.websocket import WebsocketClient def test_get_error_log(cached_client: Client) -> None: @@ -101,6 +103,17 @@ async def test_async_get_rendered_template(async_cached_client: Client) -> None: } +def test_websocket_get_rendered_template(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "render_template"` websocket command.""" + rendered_template = websocket_client.get_rendered_template( + 'The sun is {{ states("sun.sun").replace("_", " the ") }}.' + ) + assert rendered_template in { + "The sun is above the horizon.", + "The sun is below the horizon.", + } + + def test_check_api_config(cached_client: Client) -> None: """Tests the `POST /api/config/core/check_config` endpoint.""" assert cached_client.check_api_config() @@ -123,6 +136,12 @@ async def test_async_get_entities(async_cached_client: Client) -> None: assert "sun" in entities +def test_websocket_get_entities(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "get_entities"` websocket command.""" + entities = websocket_client.get_entities() + assert "sun" in entities + + def test_get_domains(cached_client: Client) -> None: """Tests the `GET /api/services` endpoint.""" domains = cached_client.get_domains() @@ -135,6 +154,12 @@ async def test_async_get_domains(async_cached_client: Client) -> None: assert "homeassistant" in domains +def test_websocket_get_domains(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "get_domains"` websocket command.""" + domains = websocket_client.get_domains() + assert "homeassistant" in domains + + def test_get_domain(cached_client: Client) -> None: """Tests the `GET /api/services` endpoint.""" domain = cached_client.get_domain("homeassistant") @@ -149,6 +174,13 @@ async def test_async_get_domain(async_cached_client: Client) -> None: assert domain.services +def test_websocket_get_domain(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "get_domain"` websocket command.""" + domain = websocket_client.get_domain("homeassistant") + assert domain is not None + assert domain.services + + def test_trigger_service(cached_client: Client) -> None: """Tests the `POST /api/services//` endpoint.""" notify = cached_client.get_domain("notify") @@ -157,6 +189,7 @@ def test_trigger_service(cached_client: Client) -> None: message="Your API Test Suite just said hello!", title="Test Suite Notifcation", ) + logging.info(resp) assert isinstance(resp, tuple) @@ -171,6 +204,17 @@ async def test_async_trigger_service(async_cached_client: Client) -> None: assert isinstance(resp, tuple) +def test_websocket_trigger_service(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "trigger_service"` websocket command.""" + notify = websocket_client.get_domain("notify") + assert notify is not None + resp = notify.persistent_notification( + message="Your API Test Suite just said hello!", title="Test Suite Notifcation" + ) + # Websocket API doesnt return changed states so we check for None + assert resp is None + + def test_trigger_service_with_response(cached_client: Client) -> None: """Tests the `POST /api/services//?return_response` endpoint.""" weather = cached_client.get_domain("weather") @@ -193,6 +237,20 @@ async def test_async_trigger_service_with_response(async_cached_client: Client) assert data is not None +def test_websocket_trigger_service_with_response( + websocket_client: WebsocketClient, +) -> None: + """Tests the `"type": "trigger_service_with_response"` websocket command.""" + weather = websocket_client.get_domain("weather") + assert weather is not None + data = weather.get_forecasts( + entity_id="weather.forecast_home", + type="hourly", + ) + # Websocket API doesnt return changed states so we check data is not None because we expect a response + assert data is not None + + def test_get_states(cached_client: Client) -> None: """Tests the `GET /api/states` endpoint.""" states = cached_client.get_states() @@ -207,6 +265,13 @@ async def test_async_get_states(async_cached_client: Client) -> None: assert isinstance(state, State) +def test_websocket_get_states(websocket_client: WebsocketClient) -> None: + """Tests the `"type": "get_states"` websocket command.""" + states = websocket_client.get_states() + for state in states: + assert isinstance(state, State) + + def test_get_state(cached_client: Client) -> None: """Tests the `GET /api/states/` endpoint.""" state = cached_client.get_state(entity_id="sun.sun") diff --git a/tests/test_errors.py b/tests/test_errors.py index ad84d02d..11630c4f 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -23,6 +23,8 @@ UnexpectedStatusCodeError, ) from homeassistant_api.processing import Processing +from homeassistant_api.utils import prepare_entity_id +from homeassistant_api.websocket import WebsocketClient def test_unauthorized() -> None: @@ -31,6 +33,14 @@ def test_unauthorized() -> None: pass +def test_websocket_unauthorized() -> None: + with pytest.raises(UnauthorizedError): + with WebsocketClient( + os.environ["HOMEASSISTANTAPI_WS_URL"], "lolthisisawrongtokenforsure" + ): + pass + + async def test_async_unauthorized() -> None: with pytest.raises(UnauthorizedError): async with Client( @@ -104,22 +114,22 @@ async def test_async_invalid_template(async_cached_client: Client) -> None: def test_prepare_entity_id(cached_client: Client) -> None: """Tests all cases for :py:meth:`Client.prepare_entity_id`.""" - assert cached_client.prepare_entity_id(group_id="person", slug="me") == "person.me" - assert cached_client.prepare_entity_id(entity_id="person.me") == "person.me" - assert "person.you" == cached_client.prepare_entity_id( + assert prepare_entity_id(group_id="person", slug="me") == "person.me" + assert prepare_entity_id(entity_id="person.me") == "person.me" + assert "person.you" == prepare_entity_id( group_id="person", entity_id="person.you", ) - assert "person.you" == cached_client.prepare_entity_id( + assert "person.you" == prepare_entity_id( slug="me", entity_id="person.you", ) with pytest.raises(ValueError): - cached_client.prepare_entity_id(group_id="person") # No slug + prepare_entity_id(group_id="person") # No slug with pytest.raises(ValueError): - cached_client.prepare_entity_id(slug="me") # No group + prepare_entity_id(slug="me") # No group with pytest.raises(ValueError): - cached_client.prepare_entity_id() # No entity_id + prepare_entity_id() # No entity_id def make_response( @@ -206,6 +216,6 @@ def test_exception_unexpected_status_code() -> None: Processing(make_response(0, "", {})).process() -def test_unkown_scheme(cached_client: Client) -> None: +def test_unkown_scheme() -> None: with pytest.raises(ValueError): - Client("ftp://example.com", "token") \ No newline at end of file + Client("ftp://example.com", "token") diff --git a/tests/test_events.py b/tests/test_events.py new file mode 100644 index 00000000..150065b4 --- /dev/null +++ b/tests/test_events.py @@ -0,0 +1,30 @@ +from datetime import datetime + +import pytest + +from homeassistant_api.websocket import WebsocketClient + + +def test_subscribe_events(websocket_client: WebsocketClient) -> None: + with websocket_client.subscribe_events("test_event") as events: + websocket_client.fire_event( + "test_event", message="Triggered by websocket client" + ) + for _, event in zip(range(1), events): + assert event.origin == "LOCAL" + assert event.event_type == "test_event" + assert event.data["message"] == "Triggered by websocket client" + + +def test_subscribe_trigger(websocket_client: WebsocketClient) -> None: + future = datetime.fromisoformat( + websocket_client.get_rendered_template("{{ (now() + timedelta(seconds=1)) }}") + ) + with websocket_client.subscribe_trigger( + "time", at=future.strftime("%H:%M:%S") + ) as triggers: + for _, trigger in zip(range(1), triggers): + assert trigger["trigger"]["platform"] == "time" + assert datetime.fromisoformat( + trigger["trigger"]["now"] + ).timestamp() == pytest.approx(future.timestamp(), abs=0.5)