diff --git a/brewtils/rest/publish_client.py b/brewtils/rest/publish_client.py index 9712ac09..be96fac2 100644 --- a/brewtils/rest/publish_client.py +++ b/brewtils/rest/publish_client.py @@ -158,7 +158,9 @@ def _get_parent_for_request(self): return Request(id=str(parent.id)) - def register_command(self, topic_name: str, cmd_name: str = None, cmd_func = None) -> Topic: + def register_command( + self, topic_name: str, cmd_name: str = None, cmd_func=None + ) -> Topic: """Register a command to subscribe to the topic provided. The subscriber is marked as GENERATED, and will be pruned after the system shuts down @@ -200,21 +202,25 @@ def register_command(self, topic_name: str, cmd_name: str = None, cmd_func = Non ) cmd_name = cmd_func._command.name - return self._easy_client.update_topic( topic_name=topic_name, - add=self._schema_parser.serialize_subscriber(Subscriber( - garden=brewtils.plugin.CONFIG.garden, - namespace=brewtils.plugin.CONFIG.namespace, - system=brewtils.plugin.CONFIG.name, - version=brewtils.plugin.CONFIG.version, - instance=brewtils.plugin.CONFIG.instance_name, - command=cmd_name, - subscriber_type="GENERATED", - ), to_string=False), + add=self._schema_parser.serialize_subscriber( + Subscriber( + garden=brewtils.plugin.CONFIG.garden, + namespace=brewtils.plugin.CONFIG.namespace, + system=brewtils.plugin.CONFIG.name, + version=brewtils.plugin.CONFIG.version, + instance=brewtils.plugin.CONFIG.instance_name, + command=cmd_name, + subscriber_type="GENERATED", + ), + to_string=False, + ), ) - def unregister_command(self, topic_name: str, cmd_name: str = None, cmd_func = None) -> Topic: + def unregister_command( + self, topic_name: str, cmd_name: str = None, cmd_func=None + ) -> Topic: """Unregister a command to subscribe to the topic provided. Args: @@ -256,13 +262,16 @@ def unregister_command(self, topic_name: str, cmd_name: str = None, cmd_func = N return self._easy_client.update_topic( topic_name=topic_name, - remove=self._schema_parser.serialize_subscriber(Subscriber( - garden=brewtils.plugin.CONFIG.garden, - namespace=brewtils.plugin.CONFIG.namespace, - system=brewtils.plugin.CONFIG.name, - version=brewtils.plugin.CONFIG.version, - instance=brewtils.plugin.CONFIG.instance_name, - command=cmd_name, - subscriber_type="GENERATED", - ), to_string=False), + remove=self._schema_parser.serialize_subscriber( + Subscriber( + garden=brewtils.plugin.CONFIG.garden, + namespace=brewtils.plugin.CONFIG.namespace, + system=brewtils.plugin.CONFIG.name, + version=brewtils.plugin.CONFIG.version, + instance=brewtils.plugin.CONFIG.instance_name, + command=cmd_name, + subscriber_type="GENERATED", + ), + to_string=False, + ), ) diff --git a/test/rest/publish_client_test.py b/test/rest/publish_client_test.py index 1308f1c5..0383a4f2 100644 --- a/test/rest/publish_client_test.py +++ b/test/rest/publish_client_test.py @@ -103,14 +103,14 @@ def test_register_command(self, client, easy_client): def _cmd(self, x): return x - client.register_command("topic", _cmd) + client.register_command(topic_name="topic", cmd_func=_cmd) easy_client.update_topic.assert_called() def test_register_command_string(self, client, easy_client): self.setup_config() - client.register_command("topic", "command") + client.register_command(topic_name="topic", cmd_name="command") easy_client.update_topic.assert_called() @@ -121,7 +121,7 @@ def test_unregister_command(self, client, easy_client): def _cmd(self, x): return x - client.unregister_command("topic", _cmd) + client.unregister_command(topic_name="topic", cmd_func=_cmd) easy_client.update_topic.assert_called() @@ -139,7 +139,7 @@ def _cmd(self, x): return x with pytest.raises(BrewtilsException): - client.register_command("topic", _cmd) + client.register_command(topic_name="topic", cmd_func=_cmd) def test_unregister_command_non_annotated(self, client): self.setup_config() @@ -148,7 +148,7 @@ def _cmd(self, x): return x with pytest.raises(BrewtilsException): - client.unregister_command("topic", _cmd) + client.unregister_command(topic_name="topic", cmd_func=_cmd) def test_register_command_no_config(self, client): @@ -157,7 +157,7 @@ def _cmd(self, x): return x with pytest.raises(BrewtilsException): - client.register_command("topic", _cmd) + client.register_command(topic_name="topic", cmd_func=_cmd) def test_unregister_command_no_config(self, client): @@ -166,4 +166,4 @@ def _cmd(self, x): return x with pytest.raises(BrewtilsException): - client.unregister_command("topic", _cmd) + client.unregister_command(topic_name="topic", cmd_func=_cmd)