From e0f4e3a145c60e500bde930a4caa57481c8711c8 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Sat, 21 Dec 2024 21:07:18 -0600
Subject: [PATCH 01/17] Add ws trigger_service

---
 homeassistant_api/websocket.py | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 0fb5f5c6..36b6fbf3 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -34,9 +34,9 @@ def get_entities(self) -> list[dict[str, str]]:
         # Note: Even though it says "get_states" this is actually comparable
         # to the `get_entities` method from the REST API clients.
         # TODO: do the same parsing logic as in the REST API client
-        return self.recv(self.send("get_states"))
+        return self.recv(self.send("get_states"))["result"]
 
-    def get_domains(self) -> list[str]:
+    def get_domains(self) -> dict[str, Domain]:
         """Get a list of (service) domains."""
         data = self.recv(self.send("get_services"))["result"]
         domains = map(
@@ -48,9 +48,29 @@ def get_domains(self) -> list[str]:
         )
         return {domain.domain_id: domain for domain in domains}
 
-    def trigger_service(self, domain: str, service: str, **service_data) -> None:
+    def trigger_service(
+        self,
+        domain: str,
+        service: str,
+        return_response: bool,  # Whether to return the response or not, no sensible default
+        entity_id: str | None = None,
+        **service_data,
+    ) -> None:
         """Trigger a service."""
-        pass
+        params = {
+            "domain": domain,
+            "service": service,
+            "service_data": service_data,
+            "return_response": return_response,
+        }
+        if entity_id is not None:
+            params["target"] = {"entity_id": entity_id}
+
+        data = self.recv(self.send("call_service", **params))
+
+        # TODO: handle data["result"]["context"]
+
+        return data["result"]["response"]
 
     def get_events(self) -> list[dict[str, str]]:
         """Get a list of events."""

From 3408df14d236f441cc79ace6f3e61a17872f16ec Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Sat, 21 Dec 2024 21:07:37 -0600
Subject: [PATCH 02/17] Fix recv event logic

---
 homeassistant_api/rawwebsocket.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py
index 4f781ab1..734d11a8 100644
--- a/homeassistant_api/rawwebsocket.py
+++ b/homeassistant_api/rawwebsocket.py
@@ -78,8 +78,9 @@ def send(self, type: str, include_id: bool = True, **data: Any) -> int:
 
         if "id" in data:
             match data["type"]:
-                case "event":
+                case "subscribe_events" | "subscribe_trigger":
                     self._event_responses[data["id"]] = []
+                    self._result_responses[data["id"]] = None
                 case "ping":
                     self._ping_responses[data["id"]] = {"start": time.perf_counter_ns()}
                 case (
@@ -135,9 +136,8 @@ def recv(self, id: int) -> dict[str, Any]:
             if self._result_responses.get(id) is not None:
                 return self._result_responses.pop(id)
             if self._event_responses.get(id, []):
-                if len(self._event_responses[id]) > 1:
+                if len(self._event_responses[id]) > 0:
                     return self._event_responses[id].pop(0)
-                return self._event_responses.pop(id)[0]
             if self._ping_responses.get(id, {}).get("end") is not None:
                 return self._ping_responses.pop(id)
 
@@ -149,10 +149,7 @@ def recv(self, id: int) -> dict[str, Any]:
                     "Received a message without an id outside the auth phase."
                 )
 
-            data = self.handle_recv(data)
-
-            if data["id"] == id:  ## we've found the message we're looking for
-                return data
+            self.handle_recv(data)
 
     def authentication_phase(self) -> dict[str, Any]:
         """Authenticate with the websocket server."""

From 14026d609373ab41eba1bde7038dd508f2d1a45b Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Mon, 30 Dec 2024 12:19:26 -0600
Subject: [PATCH 03/17] Create utils.py

---
 homeassistant_api/rawbaseclient.py | 30 ---------------------------
 homeassistant_api/utils.py         | 33 ++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 30 deletions(-)
 create mode 100644 homeassistant_api/utils.py

diff --git a/homeassistant_api/rawbaseclient.py b/homeassistant_api/rawbaseclient.py
index b5c2a8af..1b8ff006 100644
--- a/homeassistant_api/rawbaseclient.py
+++ b/homeassistant_api/rawbaseclient.py
@@ -66,36 +66,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/utils.py b/homeassistant_api/utils.py
new file mode 100644
index 00000000..afaea469
--- /dev/null
+++ b/homeassistant_api/utils.py
@@ -0,0 +1,33 @@
+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)

From 4f20f41c832e975980708991c1c672ef69171756 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:22:49 -0600
Subject: [PATCH 04/17] Add half of the tests

---
 .vscode/settings.json   |  7 +++++
 compose.yml             |  3 +-
 docs/api.rst            |  2 +-
 tests/conftest.py       | 20 +++++++++++--
 tests/test_endpoints.py | 65 +++++++++++++++++++++++++++++++++++++++++
 tests/test_errors.py    | 28 ++++++++++++------
 tests/test_events.py    |  0
 7 files changed, 110 insertions(+), 15 deletions(-)
 create mode 100644 .vscode/settings.json
 create mode 100644 tests/test_events.py

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/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/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_endpoints.py b/tests/test_endpoints.py
index 587fc140..979e9797 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."""
 
 from datetime import datetime
+import logging
 
 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/<domain>/<service>` 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/<domain>/<service>?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/<entity_id>` 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..e69de29b

From 361944b59fe44854c1c3ebbc1b9a4b5997deb61d Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:23:11 -0600
Subject: [PATCH 05/17] The rest api is not being deprecated :D

---
 homeassistant_api/client.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/homeassistant_api/client.py b/homeassistant_api/client.py
index 94c23c86..3ade46e1 100644
--- a/homeassistant_api/client.py
+++ b/homeassistant_api/client.py
@@ -3,7 +3,6 @@
 import logging
 from typing import Any
 import urllib.parse as urlparse
-import warnings
 
 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}")

From bdac620893e41aa755be2623608ae5d862d27f19 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:23:38 -0600
Subject: [PATCH 06/17] Add a message to error

---
 homeassistant_api/errors.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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):

From 7546002cfb9beeab8cac57721b67a76edd31d955 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:26:33 -0600
Subject: [PATCH 07/17] Update pydantic

---
 homeassistant_api/__init__.py       |  4 ++--
 homeassistant_api/rawasyncclient.py |  9 +++++----
 homeassistant_api/rawclient.py      | 10 ++++++----
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/homeassistant_api/__init__.py b/homeassistant_api/__init__.py
index cf924e6b..37096f2d 100644
--- a/homeassistant_api/__init__.py
+++ b/homeassistant_api/__init__.py
@@ -20,11 +20,11 @@
     "ParameterMissingError",
     "RequestError",
     "UnauthorizedError",
-    "WebSocketClient",
+    "WebsocketClient",
 )
 
 from .client import Client
-from .websocket import WebSocketClient
+from .websocket import WebsocketClient
 from .errors import (
     APIConfigurationError,
     EndpointNotFoundError,
diff --git a/homeassistant_api/rawasyncclient.py b/homeassistant_api/rawasyncclient.py
index 18c6197c..75d9392c 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:
         """
@@ -335,7 +336,7 @@ async def async_get_state(  # pylint: disable=duplicate-code
         Fetches the state of the entity specified.
         :code:`GET /api/states/<entity_id>`
         """
-        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/rawclient.py b/homeassistant_api/rawclient.py
index 2f4b92e8..48976ff0 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
@@ -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:
         """
@@ -331,7 +333,7 @@ def get_state(  # pylint: disable=duplicate-code
         Fetches the state of the entity specified.
         :code:`GET /api/states/<entity_id>`
         """
-        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))
 

From 4cb516fc50e1bab40120634412772d6110279c59 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:27:00 -0600
Subject: [PATCH 08/17] Make entity_id conditionally passed in service calls

---
 homeassistant_api/models/domains.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/homeassistant_api/models/domains.py b/homeassistant_api/models/domains.py
index 5af69f17..80520894 100644
--- a/homeassistant_api/models/domains.py
+++ b/homeassistant_api/models/domains.py
@@ -97,9 +97,11 @@ class Service(BaseModel):
     fields: Optional[Dict[str, ServiceField]] = None
 
     def trigger(
-        self, **service_data
+        self, entity_id: str | None = 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
         try:
             return self.domain._client.trigger_service_with_response(
                 self.domain.domain_id,
@@ -114,9 +116,11 @@ def trigger(
             )
 
     async def async_trigger(
-        self, **service_data
+        self, entity_id: str | None = 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
         try:
             return await self.domain._client.async_trigger_service_with_response(
                 self.domain.domain_id,
@@ -130,7 +134,7 @@ async def async_trigger(
                 **service_data,
             )
 
-    def __call__(self, **service_data) -> Union[
+    def __call__(self, entity_id: str | None = None, **service_data) -> Union[
         Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]],
         Coroutine[
             Any, Any, Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]
@@ -145,7 +149,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)

From 9a9d45bb5822f0d9fc7e3394a08fd33b8151ae46 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:28:01 -0600
Subject: [PATCH 09/17] Big commit

---
 homeassistant_api/models/websocket.py |  81 +++++++++++
 homeassistant_api/rawwebsocket.py     | 184 +++++++++++++-----------
 homeassistant_api/websocket.py        | 195 ++++++++++++++++++++++----
 3 files changed, 349 insertions(+), 111 deletions(-)
 create mode 100644 homeassistant_api/models/websocket.py

diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
new file mode 100644
index 00000000..23066d1a
--- /dev/null
+++ b/homeassistant_api/models/websocket.py
@@ -0,0 +1,81 @@
+"""A module defining the responses we expect from the websocket API."""
+
+from typing import Any, Literal, Optional
+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: Optional[float] = None  # added by the client
+    end: Optional[float] = None  # added by the client
+
+
+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[dict[str, 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"]
+    time_fired: DatetimeIsoField
+    context: Optional[Context]
+
+
+class EventResponse(BaseModel):
+    """A model to parse the response of a fired event websocket response."""
+
+    id: int
+    type: Literal["event"]
+    event: FiredEvent
diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py
index 734d11a8..44796604 100644
--- a/homeassistant_api/rawwebsocket.py
+++ b/homeassistant_api/rawwebsocket.py
@@ -1,19 +1,35 @@
 import json
 import logging
 import time
+from pydantic import ValidationError
 
 # import threading
 
 import websockets.sync.client as ws
-from typing import Any
-
-from homeassistant_api.errors import ReceivingError, ResponseError, UnauthorizedError
+from typing import Any, Optional, cast
+
+from homeassistant_api.errors import (
+    ReceivingError,
+    RequestError,
+    ResponseError,
+    UnauthorizedError,
+)
+from homeassistant_api.models.base import BaseModel
+from homeassistant_api.models.websocket import (
+    AuthInvalid,
+    AuthOk,
+    AuthRequired,
+    ErrorResponse,
+    EventResponse,
+    PingResponse,
+    ResultResponse,
+)
 
 
 logger = logging.getLogger(__name__)
 
 
-class RawWebSocketClient:
+class RawWebsocketClient:
     api_url: str
     token: str
     _conn: ws.ClientConnection
@@ -26,12 +42,15 @@ def __init__(
         self.api_url = api_url
         self.token = token
         self._conn = None
+
         self._id_counter = 0
-        self._result_responses: dict[int, dict[str, Any]] = {}  # id -> response
-        self._event_responses: dict[int, list[dict[str, Any]]] = (
+        self._result_responses: dict[int, Optional[ResultResponse]] = (
+            {}
+        )  # id -> response
+        self._event_responses: dict[int, list[EventResponse]] = (
             {}
         )  # id -> [response, ...]
-        self._ping_responses: dict[int, dict[str, float]] = {}  # id -> (sent, received)
+        self._ping_responses: dict[int, PingResponse] = {}  # id -> (sent, received)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.api_url!r})"
@@ -39,7 +58,9 @@ def __repr__(self) -> str:
     def __enter__(self):
         self._conn = ws.connect(self.api_url)
         self._conn.__enter__()
-        self.authentication_phase()
+        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):
@@ -53,15 +74,13 @@ def _request_id(self) -> int:
 
     def _send(self, data: dict[str, Any]) -> None:
         """Send a message to the websocket server."""
-        logger.info(f"Sending message: {data}")
+        logger.debug(f"Sending message: {data}")
         self._conn.send(json.dumps(data))
 
     def _recv(self) -> dict[str, Any]:
         """Receive a message from the websocket server."""
         _bytes = self._conn.recv()
-
-        # logger.info(f"Received message: {_bytes}")
-
+        logger.debug(f"Received message: {_bytes}")
         return json.loads(_bytes)
 
     def send(self, type: str, include_id: bool = True, **data: Any) -> int:
@@ -77,103 +96,104 @@ def send(self, type: str, include_id: bool = True, **data: Any) -> int:
         self._send(data)
 
         if "id" in data:
-            match data["type"]:
-                case "subscribe_events" | "subscribe_trigger":
-                    self._event_responses[data["id"]] = []
-                    self._result_responses[data["id"]] = None
-                case "ping":
-                    self._ping_responses[data["id"]] = {"start": time.perf_counter_ns()}
-                case (
-                    _
-                ):  # anything else is one-time command that returns a "type": "result" entry
-                    self._result_responses[data["id"]] = None
+            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."""
-        match data:
-            case {"type": "result", "success": False, "error": {}}:
-                raise ResponseError(data["error"].pop("message"), data["error"])
-            case {"type": "result", "success": True}:
-                # this is the expected case
-                pass
-            case {"type": "result"}:
-                raise ResponseError(
-                    "Wrongly formatted response", data
-                )  # because "type": "result" should imply a "success" key
-        return data
-
-    def handle_recv(self, data: dict[str, Any]) -> dict[str, Any]:
+        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]
+    ) -> EventResponse | ResultResponse | PingResponse:
         """Handle a received message."""
         if "id" not in data:
             raise ReceivingError(
                 "Received a message without an id outside the auth phase."
             )
-
-        match data:
-            case {"type": "pong"}:
-                logger.info("Received pong message")
-                self._ping_responses[data["id"]].update(
-                    {"end": time.perf_counter_ns(), **data}
-                )
-                data = self._ping_responses[data["id"]]
-            case {"type": "result"}:
-                logger.info("Received result message")
-                self._result_responses[data["id"]] = data
-            case {"type": "event"}:
-                logger.info("Received event message")
-                self._event_responses[data["id"]].append(data)
-            case _:
-                logger.warning(f"Received unknown message: {data}")
-
-        return self.check_success(data)
-
-    def recv(self, id: int) -> dict[str, Any]:
+        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"]] = PingResponse.model_validate(
+                {**data, "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"]["event_type"])
+            self._event_responses[data["id"]].append(EventResponse.model_validate(data))
+        else:
+            raise ReceivingError(f"Received unexpected message type: {data}")
+
+    def recv(self, id: int) -> 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 self._result_responses.pop(id)
             if self._event_responses.get(id, []):
-                if len(self._event_responses[id]) > 0:
-                    return self._event_responses[id].pop(0)
+                return self._event_responses[id].pop(0)
             if self._ping_responses.get(id, {}).get("end") is not None:
                 return self._ping_responses.pop(id)
 
             ## if not, keep receiving messages until we do
-            data = self._recv()
-
-            if "id" not in data:
-                raise ResponseError(
-                    "Received a message without an id outside the auth phase."
-                )
-
-            self.handle_recv(data)
+            self.handle_recv(self._recv())
 
-    def authentication_phase(self) -> dict[str, Any]:
+    def authentication_phase(self) -> AuthOk:
         """Authenticate with the websocket server."""
         # Capture the first message from the server saying we need to authenticate
-        welcome = self._recv()
-        logging.debug(f"Received welcome message: {welcome}")
-        if welcome["type"] != "auth_required":
-            raise ResponseError("Unexpected response during authentication")
+        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)
-        logging.debug("Sent auth message")
+        logger.debug("Sent auth message")
+
         # Check the response
-        match (resp := self._recv())["type"]:
-            case "auth_ok":
-                return None
-            case "auth_invalid":
-                raise UnauthorizedError()
-            case _:
-                raise ResponseError(
-                    "Unexpected response during authentication", resp["message"]
-                )
+        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 resp.result is None
 
     def ping_latency(self) -> float:
         """Get the latency (in milliseconds) of the connection by sending a ping message."""
-        pong = self.recv(self.send("ping"))
-        return (pong["end"] - pong["start"]) / 1_000_000
+        pong = cast(PingResponse, self.recv(self.send("ping")))
+        return (pong.end - pong.start) / 1_000_000
diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 36b6fbf3..8d6b5add 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -1,7 +1,10 @@
-from typing import Any, cast
+import contextlib
+from typing import Any, Dict, Generator, Optional, Tuple, cast
 
-from homeassistant_api.models.domains import Domain
-from .rawwebsocket import RawWebSocketClient
+from homeassistant_api.models import Domain, Entity, State, Group
+from homeassistant_api.models.states import Context
+from homeassistant_api.utils import prepare_entity_id
+from .rawwebsocket import RawWebsocketClient
 
 import urllib.parse as urlparse
 
@@ -11,7 +14,7 @@
 logger = logging.getLogger(__name__)
 
 
-class WebSocketClient(RawWebSocketClient):
+class WebsocketClient(RawWebsocketClient):
     def __init__(
         self,
         api_url: str,
@@ -22,19 +25,100 @@ def __init__(
         if parsed.scheme not in {"ws", "wss"}:
             raise ValueError(f"Unknown scheme {parsed.scheme} in {api_url}")
         super().__init__(api_url, token)
-        logger.info(f"WebSocketClient initialized with api_url: {api_url}")
+        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 first["result"] is None
+        second = self.recv(id)
+        return second["event"]["result"]
 
     def get_config(self) -> dict[str, Any]:
-        """Get the configuration."""
+        """Get the Home Assistant configuration."""
         return self.recv(self.send("get_config"))["result"]
 
-    def get_entities(self) -> list[dict[str, str]]:
-        """Get a list of entities."""
+    def get_states(self) -> Tuple[State, ...]:
+        """Get a list of states."""
+        return [
+            State.from_json(state)
+            for state in self.recv(self.send("get_states"))["result"]
+        ]
 
-        # Note: Even though it says "get_states" this is actually comparable
-        # to the `get_entities` method from the REST API clients.
-        # TODO: do the same parsing logic as in the REST API client
-        return 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
+
+    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: str | None = None,
+        slug: str | None = None,
+        entity_id: str | None = 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."""
@@ -42,17 +126,27 @@ def get_domains(self) -> dict[str, Domain]:
         domains = map(
             lambda item: Domain.from_json(
                 {"domain": item[0], "services": item[1]},
-                client=cast(WebSocketClient, self),
+                client=cast(WebsocketClient, self),
             ),
             cast(dict[str, Any], data).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,
-        return_response: bool,  # Whether to return the response or not, no sensible default
         entity_id: str | None = None,
         **service_data,
     ) -> None:
@@ -61,7 +155,7 @@ def trigger_service(
             "domain": domain,
             "service": service,
             "service_data": service_data,
-            "return_response": return_response,
+            "return_response": False,
         }
         if entity_id is not None:
             params["target"] = {"entity_id": entity_id}
@@ -70,24 +164,67 @@ def trigger_service(
 
         # TODO: handle data["result"]["context"]
 
-        return data["result"]["response"]
+        return data["result"].get(
+            "response"
+        )  # should always be None for services without a response
 
-    def get_events(self) -> list[dict[str, str]]:
-        """Get a list of events."""
-        pass
+    def trigger_service_with_response(
+        self,
+        domain: str,
+        service: str,
+        entity_id: str | None = 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}
 
-    def subscribe_event(self, event_type: str) -> None:
-        """Subscribe to an event."""
-        pass
+        data = self.recv(self.send("call_service", **params))
 
-    def unsubscribe_event(self, event_type: str) -> None:
-        """Unsubscribe from an event."""
-        pass
+        return data["result"]["response"]
 
-    def subscribe_trigger(self, entity_id: str) -> None:
-        """Subscribe to a trigger."""
+    @contextlib.contextmanager
+    def subscribe_events(self, event_type: Optional[str] = None) -> None:
+        """
+        Subscribe to all events of a certain type and calls `unsubscribe_events` when done.
+        """
+        subscription = self._subscribe_events(event_type)
+        yield 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", **params)).id
+
+    def subscribe_triggers(self, trigger: Optional[str] = None) -> None:
         pass
 
-    def unsubscribe_trigger(self, entity_id: str) -> None:
-        """Unsubscribe from a trigger."""
+    def _subscribe_triggers(self, trigger: Optional[str]) -> None:
         pass
+
+    def _wait_for(self, subscription_id: int) -> Generator[None, None, None]:
+        """
+        An iterator that waits for events of a certain type.
+        """
+        while True:
+            yield self.recv(subscription_id)
+
+    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 resp.result is None
+
+    def fire_event(self, event_type: str, **event_data) -> Context:
+        """Fire an event."""
+        params = {"event_type": event_type}
+        if event_data:
+            params["event_data"] = event_data
+        return Context.from_json(
+            self.recv(self.send("fire_event", **params))["result"]["context"]
+        )

From 577ce64a176f678f642f787e7b958eb0a8f2a7cb Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Thu, 16 Jan 2025 16:28:13 -0600
Subject: [PATCH 10/17] Add more fields to State model

---
 homeassistant_api/models/states.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

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):

From c9ba4780e507fbcf2e19f6b58973e1ceb5f16863 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 14:13:03 -0600
Subject: [PATCH 11/17] Finish events and triggers

---
 homeassistant_api/models/websocket.py | 20 +++++++++--
 homeassistant_api/rawwebsocket.py     |  2 +-
 homeassistant_api/websocket.py        | 48 +++++++++++++++++++++------
 3 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
index 23066d1a..41875896 100644
--- a/homeassistant_api/models/websocket.py
+++ b/homeassistant_api/models/websocket.py
@@ -68,9 +68,25 @@ class FiredEvent(BaseModel):
 
     event_type: str
     data: dict[str, Any]
+
     origin: Literal["LOCAL", "REMOTE"]
-    time_fired: DatetimeIsoField
+    # 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):
@@ -78,4 +94,4 @@ class EventResponse(BaseModel):
 
     id: int
     type: Literal["event"]
-    event: FiredEvent
+    event: FiredEvent | FiredTrigger | TemplateEvent
diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py
index 44796604..adb604bb 100644
--- a/homeassistant_api/rawwebsocket.py
+++ b/homeassistant_api/rawwebsocket.py
@@ -137,7 +137,7 @@ def parse_response(self, data: dict[str, Any]) -> None:
             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"]["event_type"])
+            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}")
diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 8d6b5add..1bab8083 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -3,6 +3,7 @@
 
 from homeassistant_api.models import Domain, Entity, State, Group
 from homeassistant_api.models.states import Context
+from homeassistant_api.models.websocket import EventResponse, FiredEvent, FiredTrigger, ResultResponse
 from homeassistant_api.utils import prepare_entity_id
 from .rawwebsocket import RawWebsocketClient
 
@@ -35,9 +36,10 @@ def get_rendered_template(self, template: str) -> str:
         """
         id = self.send("render_template", template=template, report_errors=True)
         first = self.recv(id)
-        assert first["result"] is None
+        assert first.result is None
         second = self.recv(id)
-        return second["event"]["result"]
+        self._unsubscribe(id)
+        return second.event.result
 
     def get_config(self) -> dict[str, Any]:
         """Get the Home Assistant configuration."""
@@ -189,12 +191,14 @@ def trigger_service_with_response(
         return data["result"]["response"]
 
     @contextlib.contextmanager
-    def subscribe_events(self, event_type: Optional[str] = None) -> None:
+    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 self._wait_for(subscription)
+        yield cast(Generator[FiredEvent, None, None], self._wait_for(subscription))
         self._unsubscribe(subscription)
 
     def _subscribe_events(self, event_type: Optional[str]) -> int:
@@ -202,23 +206,42 @@ def _subscribe_events(self, event_type: Optional[str]) -> int:
         params = {"event_type": event_type} if event_type else {}
         return self.recv(self.send("subscribe_events", **params)).id
 
-    def subscribe_triggers(self, trigger: Optional[str] = None) -> None:
-        pass
+    @contextlib.contextmanager
+    def subscribe_trigger(self, trigger: str, **trigger_fields) -> Generator[Generator[FiredTrigger, 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 cast(Generator[FiredTrigger, None, None], self._wait_for(subscription))
+        self._unsubscribe(subscription)
 
-    def _subscribe_triggers(self, trigger: Optional[str]) -> None:
-        pass
+    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[None, None, None]:
+    def _wait_for(self, subscription_id: int) -> Generator[FiredEvent | FiredTrigger, None, None]:
         """
         An iterator that waits for events of a certain type.
         """
         while True:
-            yield self.recv(subscription_id)
+            yield 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 resp.result is None
+        self._event_responses.pop(subcription_id)
 
     def fire_event(self, event_type: str, **event_data) -> Context:
         """Fire an event."""
@@ -226,5 +249,8 @@ def fire_event(self, event_type: str, **event_data) -> Context:
         if event_data:
             params["event_data"] = event_data
         return Context.from_json(
-            self.recv(self.send("fire_event", **params))["result"]["context"]
+            cast(
+                ResultResponse,
+                self.recv(self.send("fire_event", **params)),
+            ).result["context"]
         )

From dd9e5fc9fbef67d11dbc857fa0e015efd1857cdb Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 14:32:59 -0600
Subject: [PATCH 12/17] Add tests for events and triggers

---
 homeassistant_api/models/websocket.py |  2 +-
 homeassistant_api/websocket.py        | 40 +++++++++++++++++++--------
 tests/test_events.py                  | 29 +++++++++++++++++++
 3 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
index 41875896..5e1bfbdf 100644
--- a/homeassistant_api/models/websocket.py
+++ b/homeassistant_api/models/websocket.py
@@ -60,7 +60,7 @@ class ResultResponse(BaseModel):
     id: int
     success: Literal[True]
     type: Literal["result"]
-    result: Optional[dict[str, Any]]
+    result: Optional[Any]
 
 
 class FiredEvent(BaseModel):
diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 1bab8083..50e3ee5f 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -3,7 +3,12 @@
 
 from homeassistant_api.models import Domain, Entity, State, Group
 from homeassistant_api.models.states import Context
-from homeassistant_api.models.websocket import EventResponse, FiredEvent, FiredTrigger, ResultResponse
+from homeassistant_api.models.websocket import (
+    EventResponse,
+    FiredEvent,
+    FiredTrigger,
+    ResultResponse,
+)
 from homeassistant_api.utils import prepare_entity_id
 from .rawwebsocket import RawWebsocketClient
 
@@ -49,7 +54,7 @@ def get_states(self) -> Tuple[State, ...]:
         """Get a list of states."""
         return [
             State.from_json(state)
-            for state in self.recv(self.send("get_states"))["result"]
+            for state in self.recv(self.send("get_states")).result
         ]
 
     def get_state(  # pylint: disable=duplicate-code
@@ -124,13 +129,13 @@ def get_entity(
 
     def get_domains(self) -> dict[str, Domain]:
         """Get a list of (service) domains."""
-        data = self.recv(self.send("get_services"))["result"]
+        resp = self.recv(self.send("get_services"))
         domains = map(
             lambda item: Domain.from_json(
                 {"domain": item[0], "services": item[1]},
                 client=cast(WebsocketClient, self),
             ),
-            cast(dict[str, Any], data).items(),
+            cast(dict[str, Any], resp.result).items(),
         )
         return {domain.domain_id: domain for domain in domains}
 
@@ -164,10 +169,10 @@ def trigger_service(
 
         data = self.recv(self.send("call_service", **params))
 
-        # TODO: handle data["result"]["context"]
+        # TODO: handle data["result"]["context"] ?
 
-        return data["result"].get(
-            "response"
+        assert (
+            data.result.get("response") is None
         )  # should always be None for services without a response
 
     def trigger_service_with_response(
@@ -188,11 +193,12 @@ def trigger_service_with_response(
 
         data = self.recv(self.send("call_service", **params))
 
-        return data["result"]["response"]
+        return data.result["response"]
 
     @contextlib.contextmanager
     def subscribe_events(
-        self, event_type: Optional[str] = None,
+        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.
@@ -207,7 +213,9 @@ def _subscribe_events(self, event_type: Optional[str]) -> int:
         return self.recv(self.send("subscribe_events", **params)).id
 
     @contextlib.contextmanager
-    def subscribe_trigger(self, trigger: str, **trigger_fields) -> Generator[Generator[FiredTrigger, None, None], None, None]:
+    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).
@@ -219,7 +227,13 @@ def subscribe_trigger(self, trigger: str, **trigger_fields) -> Generator[Generat
         ``` -> `subscribe_trigger("state", entity_id="light.kitchen")`
         """
         subscription = self._subscribe_trigger(trigger, **trigger_fields)
-        yield cast(Generator[FiredTrigger, None, None], self._wait_for(subscription))
+        yield map(
+            lambda x: x.variables,
+            cast(
+                Generator[FiredTrigger, None, None],
+                self._wait_for(subscription),
+            ),
+        )
         self._unsubscribe(subscription)
 
     def _subscribe_trigger(self, trigger: str, **trigger_fields) -> int:
@@ -230,7 +244,9 @@ def _subscribe_trigger(self, trigger: str, **trigger_fields) -> int:
             )
         ).id
 
-    def _wait_for(self, subscription_id: int) -> Generator[FiredEvent | FiredTrigger, None, None]:
+    def _wait_for(
+        self, subscription_id: int
+    ) -> Generator[FiredEvent | FiredTrigger, None, None]:
         """
         An iterator that waits for events of a certain type.
         """
diff --git a/tests/test_events.py b/tests/test_events.py
index e69de29b..e5715c2a 100644
--- a/tests/test_events.py
+++ b/tests/test_events.py
@@ -0,0 +1,29 @@
+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)

From 8fa962adbceda990b278847a45715d331c175e63 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 14:38:18 -0600
Subject: [PATCH 13/17] Add websocket example to readme

---
 README.md | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

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(
     '<API Server URL>', # i.e. 'http://homeassistant.local:8123/api/'
     '<Your Long Lived Access-Token>'
 ) 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(
+        '<REST API Server URL>', # i.e. 'http://homeassistant.local:8123/api/'
+        '<Your Long Lived Access-Token>',
+        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(
+    '<WS API Server URL>', # i.e. 'ws://homeassistant.local:8123/api/websocket'
+    '<Your Long Lived Access-Token>'
+) 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

From 61c03f21ad110d3d30dfff4db67c67ef9c705f4e Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 15:08:17 -0600
Subject: [PATCH 14/17] Fix pinging

---
 homeassistant_api/models/websocket.py |  6 +++---
 tests/test_client.py                  | 10 +++++++++-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
index 5e1bfbdf..8bd227b1 100644
--- a/homeassistant_api/models/websocket.py
+++ b/homeassistant_api/models/websocket.py
@@ -1,10 +1,10 @@
 """A module defining the responses we expect from the websocket API."""
 
 from typing import Any, Literal, Optional
+
 from .base import BaseModel
 from .states import Context, DatetimeIsoField
 
-
 __all__ = (
     "AuthRequired",
     "AuthOk",
@@ -36,8 +36,8 @@ class PingResponse(BaseModel):
 
     id: int
     type: Literal["pong"]
-    start: Optional[float] = None  # added by the client
-    end: Optional[float] = None  # added by the client
+    start: Optional[int] = None  # added by the client, nanoseconds
+    end: Optional[int] = None  # added by the client, nanoseconds
 
 
 class Error(BaseModel):
diff --git a/tests/test_client.py b/tests/test_client.py
index 6c505b55..50e98e5d 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
\ No newline at end of file

From de6d5de401699621013090230c685c4d82306c5b Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 15:36:11 -0600
Subject: [PATCH 15/17] Lint and type check

---
 homeassistant_api/__init__.py         |  2 +-
 homeassistant_api/client.py           |  2 +-
 homeassistant_api/models/base.py      |  4 +-
 homeassistant_api/models/domains.py   | 26 +++++++--
 homeassistant_api/models/websocket.py |  2 +-
 homeassistant_api/rawbaseclient.py    |  3 +-
 homeassistant_api/rawwebsocket.py     | 35 ++++++------
 homeassistant_api/utils.py            |  1 -
 homeassistant_api/websocket.py        | 80 +++++++++++++++++----------
 tests/test_client.py                  |  2 +-
 tests/test_endpoints.py               |  2 +-
 tests/test_events.py                  |  1 +
 12 files changed, 97 insertions(+), 63 deletions(-)

diff --git a/homeassistant_api/__init__.py b/homeassistant_api/__init__.py
index 37096f2d..da79af0a 100644
--- a/homeassistant_api/__init__.py
+++ b/homeassistant_api/__init__.py
@@ -24,7 +24,6 @@
 )
 
 from .client import Client
-from .websocket import WebsocketClient
 from .errors import (
     APIConfigurationError,
     EndpointNotFoundError,
@@ -38,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 3ade46e1..c5ddd6eb 100644
--- a/homeassistant_api/client.py
+++ b/homeassistant_api/client.py
@@ -1,8 +1,8 @@
 """Module containing the primary Client class."""
 
 import logging
-from typing import Any
 import urllib.parse as urlparse
+from typing import Any
 
 from .rawasyncclient import RawAsyncClient
 from .rawclient import RawClient
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 80520894..6af8f005 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.")
@@ -98,7 +105,7 @@ class Service(BaseModel):
 
     def trigger(
         self, entity_id: str | None = None, **service_data
-    ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]]:
+    ) -> 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
@@ -121,6 +128,13 @@ async def async_trigger(
         """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,
@@ -135,7 +149,7 @@ async def async_trigger(
             )
 
     def __call__(self, entity_id: str | None = None, **service_data) -> Union[
-        Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]]],
+        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]]]
         ],
diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
index 8bd227b1..e89b0af0 100644
--- a/homeassistant_api/models/websocket.py
+++ b/homeassistant_api/models/websocket.py
@@ -36,7 +36,7 @@ class PingResponse(BaseModel):
 
     id: int
     type: Literal["pong"]
-    start: Optional[int] = None  # added by the client, nanoseconds
+    start: int  # added by the client, nanoseconds
     end: Optional[int] = None  # added by the client, nanoseconds
 
 
diff --git a/homeassistant_api/rawbaseclient.py b/homeassistant_api/rawbaseclient.py
index 1b8ff006..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
 
diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py
index adb604bb..72556992 100644
--- a/homeassistant_api/rawwebsocket.py
+++ b/homeassistant_api/rawwebsocket.py
@@ -1,12 +1,10 @@
 import json
 import logging
 import time
-from pydantic import ValidationError
-
-# import threading
+from typing import Any, Optional, cast
 
 import websockets.sync.client as ws
-from typing import Any, Optional, cast
+from pydantic import ValidationError
 
 from homeassistant_api.errors import (
     ReceivingError,
@@ -14,7 +12,6 @@
     ResponseError,
     UnauthorizedError,
 )
-from homeassistant_api.models.base import BaseModel
 from homeassistant_api.models.websocket import (
     AuthInvalid,
     AuthOk,
@@ -32,7 +29,7 @@
 class RawWebsocketClient:
     api_url: str
     token: str
-    _conn: ws.ClientConnection
+    _conn: Optional[ws.ClientConnection]
 
     def __init__(
         self,
@@ -75,12 +72,16 @@ def _request_id(self) -> int:
     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(f"Received message: {_bytes}")
+        logger.debug("Received message: %s", _bytes)
         return json.loads(_bytes)
 
     def send(self, type: str, include_id: bool = True, **data: Any) -> int:
@@ -116,9 +117,7 @@ def check_success(self, data: dict[str, Any]) -> None:
         except ValidationError:
             pass
 
-    def handle_recv(
-        self, data: dict[str, Any]
-    ) -> EventResponse | ResultResponse | PingResponse:
+    def handle_recv(self, data: dict[str, Any]) -> None:
         """Handle a received message."""
         if "id" not in data:
             raise ReceivingError(
@@ -130,9 +129,7 @@ def handle_recv(
     def parse_response(self, data: dict[str, Any]) -> None:
         if data.get("type") == "pong":
             logger.info("Received pong message")
-            self._ping_responses[data["id"]] = PingResponse.model_validate(
-                {**data, "end": time.perf_counter_ns()}
-            )
+            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)
@@ -147,11 +144,14 @@ def recv(self, id: int) -> EventResponse | ResultResponse | PingResponse:
         while True:
             ## have we received a message with the id we're looking for?
             if self._result_responses.get(id) is not None:
-                return self._result_responses.pop(id)
+                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, {}).get("end") is not None:
-                return self._ping_responses.pop(id)
+            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())
@@ -191,9 +191,10 @@ def supported_features_phase(self) -> None:
                 },
             )
         )
-        assert resp.result is None
+        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
index afaea469..ff2947ce 100644
--- a/homeassistant_api/utils.py
+++ b/homeassistant_api/utils.py
@@ -1,5 +1,4 @@
 import re
-
 from typing import Optional
 
 
diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 50e3ee5f..30deb3b3 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -1,21 +1,20 @@
 import contextlib
+import logging
+import urllib.parse as urlparse
 from typing import Any, Dict, Generator, Optional, Tuple, cast
 
-from homeassistant_api.models import Domain, Entity, State, Group
+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
-
-import urllib.parse as urlparse
-
-import logging
 
+from .rawwebsocket import RawWebsocketClient
 
 logger = logging.getLogger(__name__)
 
@@ -41,21 +40,30 @@ def get_rendered_template(self, template: str) -> str:
         """
         id = self.send("render_template", template=template, report_errors=True)
         first = self.recv(id)
-        assert first.result is None
+        assert cast(ResultResponse, first).result is None
         second = self.recv(id)
         self._unsubscribe(id)
-        return second.event.result
+        return cast(TemplateEvent, cast(EventResponse, second).event).result
 
     def get_config(self) -> dict[str, Any]:
         """Get the Home Assistant configuration."""
-        return self.recv(self.send("get_config"))["result"]
+        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 [
+        return tuple(
             State.from_json(state)
-            for state in self.recv(self.send("get_states")).result
-        ]
+            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,
@@ -78,6 +86,7 @@ def get_state(  # pylint: disable=duplicate-code
         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]:
         """
@@ -133,9 +142,9 @@ def get_domains(self) -> dict[str, Domain]:
         domains = map(
             lambda item: Domain.from_json(
                 {"domain": item[0], "services": item[1]},
-                client=cast(WebsocketClient, self),
+                client=self,
             ),
-            cast(dict[str, Any], resp.result).items(),
+            cast(dict[str, Any], cast(ResultResponse, resp).result).items(),
         )
         return {domain.domain_id: domain for domain in domains}
 
@@ -167,12 +176,16 @@ def trigger_service(
         if entity_id is not None:
             params["target"] = {"entity_id": entity_id}
 
-        data = self.recv(self.send("call_service", **params))
+        data = self.recv(self.send("call_service", include_id=True, **params))
 
         # TODO: handle data["result"]["context"] ?
 
         assert (
-            data.result.get("response") is None
+            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(
@@ -191,9 +204,9 @@ def trigger_service_with_response(
         if entity_id is not None:
             params["target"] = {"entity_id": entity_id}
 
-        data = self.recv(self.send("call_service", **params))
+        data = self.recv(self.send("call_service", include_id=True, **params))
 
-        return data.result["response"]
+        return cast(dict[str, Any], cast(ResultResponse, data).result)["response"]
 
     @contextlib.contextmanager
     def subscribe_events(
@@ -210,7 +223,7 @@ def subscribe_events(
     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", **params)).id
+        return self.recv(self.send("subscribe_events", include_id=True, **params)).id
 
     @contextlib.contextmanager
     def subscribe_trigger(
@@ -227,12 +240,12 @@ def subscribe_trigger(
         ``` -> `subscribe_trigger("state", entity_id="light.kitchen")`
         """
         subscription = self._subscribe_trigger(trigger, **trigger_fields)
-        yield map(
-            lambda x: x.variables,
-            cast(
+        yield (
+            fired_trigger.variables
+            for fired_trigger in cast(
                 Generator[FiredTrigger, None, None],
                 self._wait_for(subscription),
-            ),
+            )
         )
         self._unsubscribe(subscription)
 
@@ -251,22 +264,29 @@ def _wait_for(
         An iterator that waits for events of a certain type.
         """
         while True:
-            yield cast(EventResponse, self.recv(subscription_id)).event
+            yield cast(
+                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 resp.result is None
+        assert cast(ResultResponse, resp).result is None
         self._event_responses.pop(subcription_id)
 
-    def fire_event(self, event_type: str, **event_data) -> Context:
+    def fire_event(self, event_type: str, include_id: bool, **event_data) -> Context:
         """Fire an event."""
-        params = {"event_type": event_type}
+        params: dict[str, Any] = {"event_type": event_type}
         if event_data:
             params["event_data"] = event_data
         return Context.from_json(
             cast(
-                ResultResponse,
-                self.recv(self.send("fire_event", **params)),
-            ).result["context"]
+                dict[str, dict[str, Any]],
+                cast(
+                    ResultResponse,
+                    self.recv(self.send("fire_event", include_id=True, **params)),
+                ).result,
+            )["context"]
         )
diff --git a/tests/test_client.py b/tests/test_client.py
index 50e98e5d..105db369 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -54,4 +54,4 @@ def test_websocket_client_ping() -> None:
         os.environ["HOMEASSISTANTAPI_WS_URL"],
         os.environ["HOMEASSISTANTAPI_TOKEN"],
     ) as client:
-        assert client.ping_latency() > 0
\ No newline at end of file
+        assert client.ping_latency() > 0
diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py
index 979e9797..cbc1b01b 100644
--- a/tests/test_endpoints.py
+++ b/tests/test_endpoints.py
@@ -1,7 +1,7 @@
 """Module for making sure endpoints that should succeed, do indeed succeed."""
 
-from datetime import datetime
 import logging
+from datetime import datetime
 
 from homeassistant_api import Client
 from homeassistant_api.models.events import Event
diff --git a/tests/test_events.py b/tests/test_events.py
index e5715c2a..150065b4 100644
--- a/tests/test_events.py
+++ b/tests/test_events.py
@@ -1,6 +1,7 @@
 from datetime import datetime
 
 import pytest
+
 from homeassistant_api.websocket import WebsocketClient
 
 

From 32c1451591312c91e41da63a4e3186163a27c579 Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 15:41:40 -0600
Subject: [PATCH 16/17] Make types pre-3.10 compatible

---
 homeassistant_api/models/domains.py   | 22 +++++++++++++++-------
 homeassistant_api/models/websocket.py |  4 ++--
 homeassistant_api/rawasyncclient.py   |  6 +++---
 homeassistant_api/rawclient.py        |  8 ++++----
 homeassistant_api/rawwebsocket.py     |  4 ++--
 homeassistant_api/websocket.py        | 19 ++++++++++---------
 6 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/homeassistant_api/models/domains.py b/homeassistant_api/models/domains.py
index 6af8f005..82d68441 100644
--- a/homeassistant_api/models/domains.py
+++ b/homeassistant_api/models/domains.py
@@ -103,9 +103,12 @@ class Service(BaseModel):
     description: Optional[str] = None
     fields: Optional[Dict[str, ServiceField]] = None
 
-    def trigger(
-        self, entity_id: str | None = None, **service_data
-    ) -> Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]], dict[str, Any], None]:
+    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
@@ -123,12 +126,12 @@ def trigger(
             )
 
     async def async_trigger(
-        self, entity_id: str | None = None, **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):
@@ -148,8 +151,13 @@ async def async_trigger(
                 **service_data,
             )
 
-    def __call__(self, entity_id: str | None = None, **service_data) -> Union[
-        Union[Tuple[State, ...], Tuple[Tuple[State, ...], Dict[str, Any]], dict[str, Any], None],
+    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]]]
         ],
diff --git a/homeassistant_api/models/websocket.py b/homeassistant_api/models/websocket.py
index e89b0af0..7aae4b38 100644
--- a/homeassistant_api/models/websocket.py
+++ b/homeassistant_api/models/websocket.py
@@ -1,6 +1,6 @@
 """A module defining the responses we expect from the websocket API."""
 
-from typing import Any, Literal, Optional
+from typing import Any, Literal, Optional, Union
 
 from .base import BaseModel
 from .states import Context, DatetimeIsoField
@@ -94,4 +94,4 @@ class EventResponse(BaseModel):
 
     id: int
     type: Literal["event"]
-    event: FiredEvent | FiredTrigger | TemplateEvent
+    event: Union[FiredEvent, FiredTrigger, TemplateEvent]
diff --git a/homeassistant_api/rawasyncclient.py b/homeassistant_api/rawasyncclient.py
index 75d9392c..778e753b 100644
--- a/homeassistant_api/rawasyncclient.py
+++ b/homeassistant_api/rawasyncclient.py
@@ -233,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`.
diff --git a/homeassistant_api/rawclient.py b/homeassistant_api/rawclient.py
index 48976ff0..053bb65d 100644
--- a/homeassistant_api/rawclient.py
+++ b/homeassistant_api/rawclient.py
@@ -87,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:
@@ -228,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`.
diff --git a/homeassistant_api/rawwebsocket.py b/homeassistant_api/rawwebsocket.py
index 72556992..4f6e26ff 100644
--- a/homeassistant_api/rawwebsocket.py
+++ b/homeassistant_api/rawwebsocket.py
@@ -1,7 +1,7 @@
 import json
 import logging
 import time
-from typing import Any, Optional, cast
+from typing import Any, Optional, Union, cast
 
 import websockets.sync.client as ws
 from pydantic import ValidationError
@@ -139,7 +139,7 @@ def parse_response(self, data: dict[str, Any]) -> None:
         else:
             raise ReceivingError(f"Received unexpected message type: {data}")
 
-    def recv(self, id: int) -> EventResponse | ResultResponse | PingResponse:
+    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?
diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 30deb3b3..5c963e99 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -1,7 +1,7 @@
 import contextlib
 import logging
 import urllib.parse as urlparse
-from typing import Any, Dict, Generator, Optional, Tuple, cast
+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
@@ -105,9 +105,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`.
@@ -163,7 +163,7 @@ def trigger_service(
         self,
         domain: str,
         service: str,
-        entity_id: str | None = None,
+        entity_id: Optional[str] = None,
         **service_data,
     ) -> None:
         """Trigger a service."""
@@ -192,7 +192,7 @@ def trigger_service_with_response(
         self,
         domain: str,
         service: str,
-        entity_id: str | None = None,
+        entity_id: Optional[str] = None,
         **service_data,
     ) -> dict[str, Any]:
         params = {
@@ -259,14 +259,15 @@ def _subscribe_trigger(self, trigger: str, **trigger_fields) -> int:
 
     def _wait_for(
         self, subscription_id: int
-    ) -> Generator[FiredEvent | FiredTrigger, None, None]:
+    ) -> Generator[Union[FiredEvent, FiredTrigger], None, None]:
         """
         An iterator that waits for events of a certain type.
         """
         while True:
             yield cast(
-                FiredEvent
-                | FiredTrigger,  # we can cast this because TemplateEvent is only used for rendering templates
+                Union[
+                    FiredEvent, FiredTrigger
+                ],  # we can cast this because TemplateEvent is only used for rendering templates
                 cast(EventResponse, self.recv(subscription_id)).event,
             )
 

From 78b5385ea5804f06854587819bb9b22ed1ad85fb Mon Sep 17 00:00:00 2001
From: Nate <nlarsen23.student@gmail.com>
Date: Fri, 17 Jan 2025 15:45:17 -0600
Subject: [PATCH 17/17] Fix typo

---
 homeassistant_api/websocket.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/homeassistant_api/websocket.py b/homeassistant_api/websocket.py
index 5c963e99..8068ade2 100644
--- a/homeassistant_api/websocket.py
+++ b/homeassistant_api/websocket.py
@@ -277,7 +277,7 @@ def _unsubscribe(self, subcription_id: int) -> None:
         assert cast(ResultResponse, resp).result is None
         self._event_responses.pop(subcription_id)
 
-    def fire_event(self, event_type: str, include_id: bool, **event_data) -> Context:
+    def fire_event(self, event_type: str, **event_data) -> Context:
         """Fire an event."""
         params: dict[str, Any] = {"event_type": event_type}
         if event_data: