diff --git a/src/judas_protocol/message.py b/src/judas_protocol/message.py index 2388f14..426c13f 100644 --- a/src/judas_protocol/message.py +++ b/src/judas_protocol/message.py @@ -52,6 +52,63 @@ class Message: f", ack_required={self.ack_required})" ) + class Control: + """CONTROL category message factory methods.""" + + @staticmethod + def ack(target_id: str) -> Message: + """Create an ACK message. + + Args: + target_id (str): The ID of the message to acknowledge. + + Returns: + Message: The created ACK message. + """ + return Message( + id_=None, + category=Category.CONTROL, + action=ControlAction.ACK, + payload={"target_id": target_id}, + acknowledged=False, + ack_required=False, + ) + + @staticmethod + def hello(id_: str) -> Message: + """Create a HELLO message. + + Args: + id_ (str): The ID to include in the HELLO message. + Returns: + Message: The created HELLO message. + """ + return Message( + id_=None, + category=Category.CONTROL, + action=ControlAction.HELLO, + payload={"id": id_}, + acknowledged=False, + ack_required=True, + ) + + @staticmethod + def close() -> Message: + """Create a CLOSE message. + Prompts the recipient to close the connection and not reconnect. + + Returns: + Message: The created CLOSE message. + """ + return Message( + id_=None, + category=Category.CONTROL, + action=ControlAction.CLOSE, + payload={}, + acknowledged=False, + ack_required=True, + ) + def to_dict(self) -> dict[str, Any]: """Convert the message to a dictionary. @@ -144,57 +201,3 @@ class Message: Message: The created message. """ return cls.from_json(data.decode("utf-8")) - - @classmethod - def ack(cls, target_id: str) -> Message: - """Create an ACK message. - - Args: - target_id (str): The ID of the message to acknowledge. - - Returns: - Message: The created ACK message. - """ - return cls( - id_=None, - category=Category.CONTROL, - action=ControlAction.ACK, - payload={"target_id": target_id}, - acknowledged=False, - ack_required=False, - ) - - @classmethod - def hello(cls, id_: str) -> Message: - """Create a HELLO message. - - Args: - id_ (str): The ID to include in the HELLO message. - Returns: - Message: The created HELLO message. - """ - return cls( - id_=None, - category=Category.CONTROL, - action=ControlAction.HELLO, - payload={"id": id_}, - acknowledged=False, - ack_required=True, - ) - - @classmethod - def close(cls) -> Message: - """Create a CLOSE message. - Prompts the recipient to close the connection and not reconnect. - - Returns: - Message: The created CLOSE message. - """ - return cls( - id_=None, - category=Category.CONTROL, - action=ControlAction.CLOSE, - payload={}, - acknowledged=False, - ack_required=True, - )