diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 97f91504..a1d4b1bd 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,6 +1,12 @@ Brewtils Changelog ================== +3.27.2 +------ +TBD + +- Expanded PublishClient to support Registering and Unregistering commands to Topics after a plugin has been initialized + 3.27.1 ------ 8/21/24 diff --git a/brewtils/rest/publish_client.py b/brewtils/rest/publish_client.py index 859652e6..2c44eb9d 100644 --- a/brewtils/rest/publish_client.py +++ b/brewtils/rest/publish_client.py @@ -3,7 +3,7 @@ import brewtils.plugin from brewtils.errors import BrewtilsException -from brewtils.models import Event, Events, Request +from brewtils.models import Event, Events, Request, Subscriber, Topic from brewtils.rest.easy_client import EasyClient @@ -155,3 +155,113 @@ def _get_parent_for_request(self): return None return Request(id=str(parent.id)) + + def register_command(self, topic: str, cmd) -> 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 + + Args: + topic (str): Topic for the command to subscribe to + cmd (str, function): Command to register + + Raises: + BrewtilsException: If function is provided, it must be an annotated + function. Only supports running plugins + + Returns: + Topic: Updated Topic Model + """ + + if not ( + brewtils.plugin.CONFIG.garden + and brewtils.plugin.CONFIG.name + and brewtils.plugin.CONFIG.version + and brewtils.plugin.CONFIG.instance_name + and brewtils.plugin.CONFIG.namespace + ): + raise BrewtilsException( + ( + "Unable to identify Configuration for Plugin, " + "only running Plugins can register commands" + ) + ) + + if isinstance(cmd, str): + command_name = cmd + else: + if not hasattr(cmd, "_command"): + raise BrewtilsException( + ( + "Attempted to register command " + f"{getattr(cmd, '__name__', 'MISSING FUNC NAME')} " + "that is not an annotated command" + ) + ) + command_name = cmd._command.name + + return self._easy_client.update_topic( + topic_name=topic, + add=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=command_name, + subscriber_type="GENERATED", + ), + ) + + def unregister_command(self, topic: str, cmd) -> Topic: + """Unregister a command to subscribe to the topic provided. + + Args: + topic (str): Topic for the command to subscribe to + cmd (str, function): Command to unregister + + Raises: + BrewtilsException: If function is provided, it must be + an annotated function. Only supports running plugins + + Returns: + Topic: Updated Topic Model + """ + if not ( + brewtils.plugin.CONFIG.garden + and brewtils.plugin.CONFIG.name + and brewtils.plugin.CONFIG.version + and brewtils.plugin.CONFIG.instance_name + and brewtils.plugin.CONFIG.namespace + ): + raise BrewtilsException( + ( + "Unable to identify Configuration for Plugin, only " + "running Plugins can unregister commands" + ) + ) + + if isinstance(cmd, str): + command_name = cmd + else: + if not hasattr(cmd, "_command"): + raise BrewtilsException( + ( + "Attempted to register command " + f"{getattr(cmd, '__name__', 'MISSING FUNC NAME')} " + "that is not an annotated command" + ) + ) + command_name = cmd._command.name + + return self._easy_client.update_topic( + topic_name=topic, + remove=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=command_name, + subscriber_type="GENERATED", + ), + ) diff --git a/test/rest/publish_client_test.py b/test/rest/publish_client_test.py index 13ecc2d3..1308f1c5 100644 --- a/test/rest/publish_client_test.py +++ b/test/rest/publish_client_test.py @@ -2,6 +2,7 @@ from mock import Mock import brewtils.rest +from brewtils.decorators import command from brewtils.errors import BrewtilsException from brewtils.models import Event, Events, Request, System from brewtils.rest.publish_client import PublishClient @@ -12,6 +13,7 @@ def easy_client(monkeypatch): mock = Mock(name="easy_client") mock.publish_event.return_value = True + mock.update_topic.return_value = True monkeypatch.setattr( brewtils.rest.publish_client, "EasyClient", Mock(return_value=mock) @@ -93,3 +95,75 @@ def test_verify_generated_request(self, client, easy_client): assert SchemaParser.serialize_event( called_event ) == SchemaParser.serialize_event(event) + + def test_register_command(self, client, easy_client): + self.setup_config() + + @command + def _cmd(self, x): + return x + + client.register_command("topic", _cmd) + + easy_client.update_topic.assert_called() + + def test_register_command_string(self, client, easy_client): + self.setup_config() + + client.register_command("topic", "command") + + easy_client.update_topic.assert_called() + + def test_unregister_command(self, client, easy_client): + self.setup_config() + + @command + def _cmd(self, x): + return x + + client.unregister_command("topic", _cmd) + + easy_client.update_topic.assert_called() + + def test_unregister_command_string(self, client, easy_client): + self.setup_config() + + client.unregister_command("topic", "command") + + easy_client.update_topic.assert_called() + + def test_register_command_non_annotated(self, client): + self.setup_config() + + def _cmd(self, x): + return x + + with pytest.raises(BrewtilsException): + client.register_command("topic", _cmd) + + def test_unregister_command_non_annotated(self, client): + self.setup_config() + + def _cmd(self, x): + return x + + with pytest.raises(BrewtilsException): + client.unregister_command("topic", _cmd) + + def test_register_command_no_config(self, client): + + @command + def _cmd(self, x): + return x + + with pytest.raises(BrewtilsException): + client.register_command("topic", _cmd) + + def test_unregister_command_no_config(self, client): + + @command + def _cmd(self, x): + return x + + with pytest.raises(BrewtilsException): + client.unregister_command("topic", _cmd)