Merge pull request #502 from beer-garden/publish_register_cmd

Added Register and Un Register commands for topics

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

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",
+            ),
+        )

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)