Merge pull request #512 from beer-garden/topic_add_sub_patch

Refactor registration of commands
beer-garden · Oct 9, 2024 · 7e3aacb · 7e3aacb
2 parents 77ebeab + de8a679
commit 7e3aacb
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 43 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -11,6 +11,7 @@ TBD
 - Updated SystemClient to utilize the local Garden name for default Namespace if none can be determined
 - Updated default Garden version to `UNKNOWN`
 - Updated `get_current_request_read_only` to support sub threads calling where current_request is not populated
+- Refactored Publishclient for Registering and Unregistering commands input values
 
 3.27.2
 ------

diff --git a/brewtils/rest/publish_client.py b/brewtils/rest/publish_client.py
@@ -5,6 +5,7 @@
 from brewtils.errors import BrewtilsException
 from brewtils.models import Event, Events, Request, Subscriber, Topic
 from brewtils.rest.easy_client import EasyClient
+from brewtils.schema_parser import SchemaParser
 
 
 class PublishClient(object):
@@ -64,6 +65,7 @@ class PublishClient(object):
     def __init__(self, *args, **kwargs):
         self._logger = logging.getLogger(__name__)
         self._easy_client = EasyClient(*args, **kwargs)
+        self._schema_parser = SchemaParser()
 
     def publish(
         self,
@@ -156,13 +158,16 @@ def _get_parent_for_request(self):
 
         return Request(id=str(parent.id))
 
-    def register_command(self, topic: str, cmd) -> 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
 
         Args:
-            topic (str): Topic for the command to subscribe to
-            cmd (str, function): Command to register
+            topic_name (str): Topic for the command to subscribe to
+            cmd_name (str): Command to register
+            cmd_func (function): Command to register
 
         Raises:
             BrewtilsException: If function is provided, it must be an annotated
@@ -186,38 +191,42 @@ def register_command(self, topic: str, cmd) -> Topic:
                 )
             )
 
-        if isinstance(cmd, str):
-            command_name = cmd
-        else:
-            if not hasattr(cmd, "_command"):
+        if not cmd_name:
+            if not hasattr(cmd_func, "_command"):
                 raise BrewtilsException(
                     (
                         "Attempted to register command "
-                        f"{getattr(cmd, '__name__', 'MISSING FUNC NAME')} "
+                        f"{getattr(cmd_func, '__name__', 'MISSING FUNC NAME')} "
                         "that is not an annotated command"
                     )
                 )
-            command_name = cmd._command.name
+            cmd_name = cmd_func._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",
+            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,
             ),
         )
 
-    def unregister_command(self, topic: str, cmd) -> 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:
-            topic (str): Topic for the command to subscribe to
-            cmd (str, function): Command to unregister
+            topic_name (str): Topic for the command to subscribe to
+            cmd_name (str): Command to unregister
+            cmd_func (function): Command to unregister
 
         Raises:
             BrewtilsException: If function is provided, it must be
@@ -240,28 +249,29 @@ def unregister_command(self, topic: str, cmd) -> Topic:
                 )
             )
 
-        if isinstance(cmd, str):
-            command_name = cmd
-        else:
-            if not hasattr(cmd, "_command"):
+        if not cmd_name:
+            if not hasattr(cmd_func, "_command"):
                 raise BrewtilsException(
                     (
                         "Attempted to register command "
-                        f"{getattr(cmd, '__name__', 'MISSING FUNC NAME')} "
+                        f"{getattr(cmd_func, '__name__', 'MISSING FUNC NAME')} "
                         "that is not an annotated command"
                     )
                 )
-            command_name = cmd._command.name
+            cmd_name = cmd_func._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",
+            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,
             ),
         )
diff --git 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)