Add some entity subscription helper methods

rustplus/remote/websocket/ws.py

@@ -122,7 +122,7 @@ async def run(self) -> None:
 
     async def send_and_get(self, request: AppRequest) -> AppMessage:
         await self.send_message(request)
-        return await self.get_response(request)
+        return await self.get_response(request.seq)
 
     async def send_message(
             self, request: AppRequest, ignore_response: bool = False
@@ -141,10 +141,10 @@ async def send_message(
         except Exception as err:
             self.logger.warning("WebSocket connection error: %s", err)
 
-    async def get_response(self, request: AppRequest) -> Union[AppMessage, None]:
+    async def get_response(self, seq: int) -> Union[AppMessage, None]:
 
-        response = await self.responses[request.seq].wait(timeout=self.RESPONSE_TIMEOUT)
-        del self.responses[request.seq]
+        response = await self.responses[seq].wait(timeout=self.RESPONSE_TIMEOUT)
+        del self.responses[seq]
 
         return response
 

rustplus/rust_api.py

@@ -17,7 +17,7 @@
     AppSetEntityValue,
     AppPromoteToLeader,
     AppCameraSubscribe,
-    AppMapMonument,
+    AppMapMonument, AppFlag,
 )
 from .remote.websocket import RustWebsocket
 from .structs import (
@@ -381,7 +381,7 @@ async def get_entity_info(self, eid: int = None) -> Union[RustEntityInfo, None]:
 
         return RustEntityInfo(response.response.entity_info)
 
-    async def set_entity_value(self, eid: int = None, value: bool = False) -> None:
+    async def set_entity_value(self, eid: int, value: bool = False) -> None:
         """
         Turns off a given smart switch by entity ID
 
@@ -397,6 +397,39 @@ async def set_entity_value(self, eid: int = None, value: bool = False) -> None:
 
         await self.ws.send_message(packet, True)
 
+    async def set_subscription_to_entity(self, eid: int, value: bool = True) -> None:
+        """
+        Subscribes to an entity for events
+
+        :param eid: The Entities ID
+        :param value: The value to set the subscription to
+        :return None:
+        """
+        packet = await self._generate_request()
+        flag = AppFlag()
+        flag.value = value
+        packet.set_subscription = flag
+        packet.entity_id = eid
+
+        await self.ws.send_message(packet, True)
+
+    async def check_subscription_to_entity(self, eid: int) -> Union[bool, None]:
+        """
+        Checks if you are subscribed to an entity
+
+        :param eid: The Entities ID
+        :return bool: If you are subscribed
+        """
+        packet = await self._generate_request()
+        packet.check_subscription = AppEmpty()
+        packet.entity_id = eid
+
+        response = await self.ws.send_and_get(packet)
+        if response is None:
+            return None
+
+        return response.response.flag.value
+
     async def promote_to_team_leader(self, steamid: int = None) -> None:
         """
         Promotes a given user to the team leader by their 64-bit Steam ID