Add lock entity (#158)

aioesphomeapi/api.proto

@@ -42,6 +42,7 @@ service APIConnection {
   rpc select_command (SelectCommandRequest) returns (void) {}
   rpc siren_command (SirenCommandRequest) returns (void) {}
   rpc button_command (ButtonCommandRequest) returns (void) {}
+  rpc lock_command (LockCommandRequest) returns (void) {}
 }
 
 
@@ -986,6 +987,58 @@ message SirenCommandRequest {
   float volume = 9;
 }
 
+// ==================== LOCK ====================
+enum LockState {
+  LOCK_STATE_NONE = 0;
+  LOCK_STATE_LOCKED = 1;
+  LOCK_STATE_UNLOCKED = 2;
+  LOCK_STATE_JAMMED = 3;
+  LOCK_STATE_LOCKING = 4;
+  LOCK_STATE_UNLOCKING = 5;
+}
+enum LockCommand  {
+  LOCK_UNLOCK = 0;
+  LOCK_LOCK = 1;
+  LOCK_OPEN = 2;
+}
+message ListEntitiesLockResponse {
+  option (id) = 58;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_LOCK";
+
+  string object_id = 1;
+  fixed32 key = 2;
+  string name = 3;
+  string unique_id = 4;
+
+  string icon = 5;
+  bool disabled_by_default = 6;
+  EntityCategory entity_category = 7;
+  bool assumed_state = 8;
+  
+  bool supports_open = 9;
+  bool requires_code = 10;
+  string code_format = 11;
+}
+message LockStateResponse {
+  option (id) = 59;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_LOCK";
+  option (no_delay) = true;
+  fixed32 key = 1;
+  LockState  state = 2;
+}
+message LockCommandRequest {
+  option (id) = 60;
+  option (source) = SOURCE_CLIENT;
+  option (ifdef) = "USE_LOCK";
+  option (no_delay) = true;
+  fixed32 key = 1;
+  LockCommand command = 2;
+  bool has_code = 3;
+  string code = 4;
+}
+
 // ==================== BUTTON ====================
 message ListEntitiesButtonResponse {
   option (id) = 61;

aioesphomeapi/client.py

@@ -41,6 +41,7 @@ from .api_pb2 import (  # type: ignore
     ListEntitiesDoneResponse,
     ListEntitiesFanResponse,
     ListEntitiesLightResponse,
+    ListEntitiesLockResponse,
     ListEntitiesNumberResponse,
     ListEntitiesRequest,
     ListEntitiesSelectResponse,
@@ -49,6 +50,8 @@ from .api_pb2 import (  # type: ignore
     ListEntitiesSirenResponse,
     ListEntitiesSwitchResponse,
     ListEntitiesTextSensorResponse,
+    LockCommandRequest,
+    LockStateResponse,
     NumberCommandRequest,
     NumberStateResponse,
     SelectCommandRequest,
@@ -95,6 +98,9 @@ from .model import (
     LegacyCoverCommand,
     LightInfo,
     LightState,
+    LockCommand,
+    LockEntityState,
+    LockInfo,
     LogLevel,
     NumberInfo,
     NumberState,
@@ -233,6 +239,7 @@ class APIClient:
             ListEntitiesServicesResponse: None,
             ListEntitiesCameraResponse: CameraInfo,
             ListEntitiesClimateResponse: ClimateInfo,
+            ListEntitiesLockResponse: LockInfo,
         }
 
         def do_append(msg: message.Message) -> bool:
@@ -276,6 +283,7 @@ class APIClient:
             SwitchStateResponse: SwitchState,
             TextSensorStateResponse: TextSensorState,
             ClimateStateResponse: ClimateState,
+            LockStateResponse: LockEntityState,
         }
 
         image_stream: Dict[int, bytes] = {}
@@ -606,6 +614,22 @@ class APIClient:
         assert self._connection is not None
         await self._connection.send_message(req)
 
+    async def lock_command(
+        self,
+        key: int,
+        command: LockCommand,
+        code: Optional[str] = None,
+    ) -> None:
+        self._check_authenticated()
+
+        req = LockCommandRequest()
+        req.key = key
+        req.command = command
+        if code is not None:
+            req.code = code
+        assert self._connection is not None
+        await self._connection.send_message(req)
+
     async def execute_service(
         self, service: UserService, data: ExecuteServiceDataType
     ) -> None:

aioesphomeapi/core.py

@@ -32,6 +32,7 @@ from .api_pb2 import (  # type: ignore
     ListEntitiesDoneResponse,
     ListEntitiesFanResponse,
     ListEntitiesLightResponse,
+    ListEntitiesLockResponse,
     ListEntitiesNumberResponse,
     ListEntitiesRequest,
     ListEntitiesSelectResponse,
@@ -40,6 +41,8 @@ from .api_pb2 import (  # type: ignore
     ListEntitiesSirenResponse,
     ListEntitiesSwitchResponse,
     ListEntitiesTextSensorResponse,
+    LockCommandRequest,
+    LockStateResponse,
     NumberCommandRequest,
     NumberStateResponse,
     PingRequest,
@@ -167,6 +170,9 @@ MESSAGE_TYPE_TO_PROTO = {
     55: ListEntitiesSirenResponse,
     56: SirenStateResponse,
     57: SirenCommandRequest,
+    58: ListEntitiesLockResponse,
+    59: LockStateResponse,
+    60: LockCommandRequest,
     61: ListEntitiesButtonResponse,
     62: ButtonCommandRequest,
 }

aioesphomeapi/model.py

@@ -602,6 +602,38 @@ class ButtonInfo(EntityInfo):
     device_class: str = ""
 
 
+# ==================== LOCK ====================
+class LockState(APIIntEnum):
+    NONE = 0
+    LOCKED = 1
+    UNLOCKED = 3
+    JAMMED = 3
+    LOCKING = 4
+    UNLOCKING = 5
+
+
+class LockCommand(APIIntEnum):
+    UNLOCK = 0
+    LOCK = 1
+    OPEN = 2
+
+
+@dataclass(frozen=True)
+class LockInfo(EntityInfo):
+    supports_open: bool = False
+    assumed_state: bool = False
+
+    requires_code: bool = False
+    code_format: str = ""
+
+
+@dataclass(frozen=True)
+class LockEntityState(EntityState):
+    state: Optional[LockState] = converter_field(
+        default=LockState.NONE, converter=LockState.convert
+    )
+
+
 # ==================== INFO MAP ====================
 
 COMPONENT_TYPE_TO_INFO: Dict[str, Type[EntityInfo]] = {
@@ -618,6 +650,7 @@ COMPONENT_TYPE_TO_INFO: Dict[str, Type[EntityInfo]] = {
     "select": SelectInfo,
     "siren": SirenInfo,
     "button": ButtonInfo,
+    "lock": LockInfo,
 }
 
 

tests/test_client.py

@@ -16,6 +16,7 @@ from aioesphomeapi.api_pb2 import (
     ListEntitiesBinarySensorResponse,
     ListEntitiesDoneResponse,
     ListEntitiesServicesResponse,
+    LockCommandRequest,
     NumberCommandRequest,
     SelectCommandRequest,
     SwitchCommandRequest,
@@ -33,6 +34,7 @@ from aioesphomeapi.model import (
     FanDirection,
     FanSpeed,
     LegacyCoverCommand,
+    LockCommand,
     UserService,
     UserServiceArg,
     UserServiceArgType,
@@ -356,6 +358,25 @@ async def test_number_command(auth_client, cmd, req):
     send.assert_called_once_with(NumberCommandRequest(**req))
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "cmd, req",
+    [
+        (dict(key=1, command=LockCommand.LOCK), dict(key=1, command=LockCommand.LOCK)),
+        (
+            dict(key=1, command=LockCommand.UNLOCK),
+            dict(key=1, command=LockCommand.UNLOCK),
+        ),
+        (dict(key=1, command=LockCommand.OPEN), dict(key=1, command=LockCommand.OPEN)),
+    ],
+)
+async def test_lock_command(auth_client, cmd, req):
+    send = patch_send(auth_client)
+
+    await auth_client.lock_command(**cmd)
+    send.assert_called_once_with(LockCommandRequest(**req))
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "cmd, req",

tests/test_model.py

@@ -18,6 +18,7 @@ from aioesphomeapi.api_pb2 import (
     ListEntitiesCoverResponse,
     ListEntitiesFanResponse,
     ListEntitiesLightResponse,
+    ListEntitiesLockResponse,
     ListEntitiesNumberResponse,
     ListEntitiesSelectResponse,
     ListEntitiesSensorResponse,
@@ -25,6 +26,7 @@ from aioesphomeapi.api_pb2 import (
     ListEntitiesServicesResponse,
     ListEntitiesSwitchResponse,
     ListEntitiesTextSensorResponse,
+    LockStateResponse,
     NumberStateResponse,
     SelectStateResponse,
     SensorStateResponse,
@@ -51,6 +53,8 @@ from aioesphomeapi.model import (
     LegacyCoverState,
     LightInfo,
     LightState,
+    LockEntityState,
+    LockInfo,
     NumberInfo,
     NumberState,
     SelectInfo,
@@ -220,6 +224,8 @@ def test_api_version_ord():
         (UserServiceArg, ListEntitiesServicesArgument),
         (UserService, ListEntitiesServicesResponse),
         (ButtonInfo, ListEntitiesButtonResponse),
+        (LockInfo, ListEntitiesLockResponse),
+        (LockEntityState, LockStateResponse),
     ],
 )
 def test_basic_pb_conversions(model, pb):