Skip to content

Commit c2b5d04

Browse files
dmulcaheydmulcahey
committed
streamline
1 parent 352d47b commit c2b5d04

File tree

3 files changed

+48
-71
lines changed

3 files changed

+48
-71
lines changed

‎zha/application/platforms/websocket_api.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@
44

55
import inspect
66
import logging
7-
from typing import TYPE_CHECKING, Any, Literal
7+
from typing import TYPE_CHECKING, Literal
88

99
from zigpy.types.named import EUI64
1010

1111
from zha.application import Platform
12-
from zha.websocket.const import ATTR_UNIQUE_ID, IEEE, APICommands
12+
from zha.websocket.const import APICommands
1313
from zha.websocket.server.api import decorators, register_api_command
1414
from zha.websocket.server.api.model import WebSocketCommand
1515

@@ -74,13 +74,7 @@ async def execute_platform_entity_command(
7474
client.send_result_error(command, "PLATFORM_ENTITY_ACTION_ERROR", str(err))
7575
return
7676

77-
result: dict[str, Any] = {}
78-
if command.ieee:
79-
result[IEEE] = str(command.ieee)
80-
else:
81-
result["group_id"] = command.group_id
82-
result[ATTR_UNIQUE_ID] = command.unique_id
83-
client.send_result_success(command, result)
77+
client.send_result_success(command)
8478

8579

8680
class PlatformEntityRefreshStateCommand(PlatformEntityCommand):

‎zha/application/websocket_api.py‎

Lines changed: 39 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from pydantic import Field
1010
from zigpy.types.named import EUI64
1111

12-
from zha.websocket.const import GROUPS, APICommands
12+
from zha.websocket.const import DEVICE, DEVICES, GROUPS, APICommands
1313
from zha.websocket.server.api import decorators, register_api_command
1414
from zha.websocket.server.api.model import (
1515
GetApplicationStateResponse,
@@ -101,16 +101,16 @@ async def get_devices(
101101
) -> None:
102102
"""Get Zigbee devices."""
103103
try:
104-
response = GetDevicesResponse(
105-
success=True,
106-
devices={
107-
ieee: device.extended_device_info
108-
for ieee, device in gateway.devices.items()
104+
client.send_result_success(
105+
command,
106+
data={
107+
DEVICES: {
108+
ieee: device.extended_device_info
109+
for ieee, device in gateway.devices.items()
110+
}
109111
},
110-
message_id=command.message_id,
112+
response_type=GetDevicesResponse,
111113
)
112-
_LOGGER.info("response: %s", response)
113-
client.send_result_success(command, response)
114114
except Exception as e:
115115
_LOGGER.exception("Error getting devices", exc_info=e)
116116
client.send_result_error(command, "Error getting devices", str(e))
@@ -154,12 +154,7 @@ async def get_groups(
154154
) # maybe we should change the group_id type...
155155
_LOGGER.info("groups: %s", groups)
156156
client.send_result_success(
157-
command,
158-
GroupsResponse(
159-
**command.model_dump(exclude="model_class_name"),
160-
groups=groups,
161-
success=True,
162-
),
157+
command, data={GROUPS: groups}, response_type=GroupsResponse
163158
)
164159

165160

@@ -179,13 +174,7 @@ async def permit_joining(
179174
"""Permit joining devices to the Zigbee network."""
180175
# TODO add permit with code support
181176
await gateway.application_controller.permit(command.duration, command.ieee)
182-
response = PermitJoiningResponse(
183-
**command.model_dump(exclude="model_class_name"),
184-
success=True,
185-
duration=command.duration,
186-
ieee=command.ieee,
187-
)
188-
client.send_result_success(command, response)
177+
client.send_result_success(command, response_type=PermitJoiningResponse)
189178

190179

191180
class RemoveDeviceCommand(WebSocketCommand):
@@ -256,22 +245,22 @@ async def read_cluster_attributes(
256245
attributes, allow_cache=False, only_cache=False, manufacturer=manufacturer
257246
)
258247

259-
response = ReadClusterAttributesResponse(
260-
message_id=command.message_id,
261-
success=True,
262-
device=device.extended_device_info,
263-
cluster={
248+
data = {
249+
DEVICE: device.extended_device_info,
250+
"cluster": {
264251
"id": cluster.cluster_id,
265252
"name": cluster.name,
266253
"type": cluster.cluster_type,
267254
"endpoint_id": cluster.endpoint.endpoint_id,
268255
"endpoint_attribute": cluster.ep_attribute,
269256
},
270-
manufacturer_code=manufacturer,
271-
succeeded=success,
272-
failed=failure,
257+
"succeeded": success,
258+
"failed": failure,
259+
}
260+
261+
client.send_result_success(
262+
command, data=data, response_type=ReadClusterAttributesResponse
273263
)
274-
client.send_result_success(command, response)
275264

276265

277266
class WriteClusterAttributeCommand(WebSocketCommand):
@@ -332,24 +321,24 @@ async def write_cluster_attribute(
332321
manufacturer=manufacturer,
333322
)
334323

335-
api_response = WriteClusterAttributeResponse(
336-
message_id=command.message_id,
337-
success=True,
338-
device=device.extended_device_info,
339-
cluster={
324+
data = {
325+
DEVICE: device.extended_device_info,
326+
"cluster": {
340327
"id": cluster.cluster_id,
341328
"name": cluster.name,
342329
"type": cluster.cluster_type,
343330
"endpoint_id": cluster.endpoint.endpoint_id,
344331
"endpoint_attribute": cluster.ep_attribute,
345332
},
346-
manufacturer_code=manufacturer,
347-
response={
333+
"response": {
348334
"attribute": attribute,
349335
"status": response[0][0].status.name, # type: ignore
350336
}, # TODO there has to be a better way to do this
337+
}
338+
339+
client.send_result_success(
340+
command, data=data, response_type=WriteClusterAttributeResponse
351341
)
352-
client.send_result_success(command, api_response)
353342

354343

355344
class CreateGroupCommand(WebSocketCommand):
@@ -371,12 +360,9 @@ async def create_group(
371360
members = command.members
372361
group_id = command.group_id
373362
group: Group = await gateway.async_create_zigpy_group(group_name, members, group_id)
374-
response = UpdateGroupResponse(
375-
**command.model_dump(exclude="model_class_name"),
376-
group=group.info_object,
377-
success=True,
363+
client.send_result_success(
364+
command, data={GROUP: group.info_object}, response_type=UpdateGroupResponse
378365
)
379-
client.send_result_success(command, response)
380366

381367

382368
class RemoveGroupsCommand(WebSocketCommand):
@@ -405,7 +391,9 @@ async def remove_groups(
405391
for group_id, group in gateway.groups.items():
406392
groups[int(group_id)] = group.info_object
407393
_LOGGER.info("groups: %s", groups)
408-
client.send_result_success(command, {GROUPS: groups})
394+
client.send_result_success(
395+
command, data={GROUPS: groups}, response_type=GroupsResponse
396+
)
409397

410398

411399
class AddGroupMembersCommand(WebSocketCommand):
@@ -434,12 +422,9 @@ async def add_group_members(
434422
if not group:
435423
client.send_result_error(command, "G1", "ZHA Group not found")
436424
return
437-
response = UpdateGroupResponse(
438-
**command.model_dump(exclude="model_class_name"),
439-
group=group.info_object,
440-
success=True,
425+
client.send_result_success(
426+
command, data={GROUP: group.info_object}, response_type=UpdateGroupResponse
441427
)
442-
client.send_result_success(command, response)
443428

444429

445430
class RemoveGroupMembersCommand(AddGroupMembersCommand):
@@ -466,12 +451,9 @@ async def remove_group_members(
466451
if not group:
467452
client.send_result_error(command, "G1", "ZHA Group not found")
468453
return
469-
response = UpdateGroupResponse(
470-
**command.model_dump(exclude="model_class_name"),
471-
group=group.info_object,
472-
success=True,
454+
client.send_result_success(
455+
command, data={GROUP: group.info_object}, response_type=UpdateGroupResponse
473456
)
474-
client.send_result_success(command, response)
475457

476458

477459
class StopServerCommand(WebSocketCommand):
@@ -505,10 +487,9 @@ async def get_application_state(
505487
) -> None:
506488
"""Get the application state."""
507489
state = gateway.application_controller.state
508-
response = GetApplicationStateResponse(
509-
success=True, message_id=command.message_id, state=state
490+
client.send_result_success(
491+
command, data={"state": state}, response_type=GetApplicationStateResponse
510492
)
511-
client.send_result_success(command, data=response)
512493

513494

514495
def load_api(gateway: WebSocketServerGateway) -> None:

‎zha/websocket/server/client.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,10 @@ def send_event(self, message: BaseEvent) -> None:
6767
self._send_data(message)
6868

6969
def send_result_success(
70-
self, command: WebSocketCommand, data: dict[str, Any] | BaseModel | None = None
70+
self,
71+
command: WebSocketCommand,
72+
data: dict[str, Any] | BaseModel | None = None,
73+
response_type: type[WebSocketCommandResponse] = WebSocketCommandResponse,
7174
) -> None:
7275
"""Send success result prompted by a client request."""
7376
if data and isinstance(data, BaseModel):
@@ -76,10 +79,9 @@ def send_result_success(
7679
if data is None:
7780
data = {}
7881
self._send_data(
79-
WebSocketCommandResponse(
82+
response_type(
83+
**command.model_dump(exclude=["model_class_name"]),
8084
success=True,
81-
message_id=command.message_id,
82-
command=command.command,
8385
**data,
8486
)
8587
)

0 commit comments

Comments
 (0)