From b84f0dd2ba5f8159fd211ef79e19ebea2f35ad91 Mon Sep 17 00:00:00 2001
From: 1maple1 <160027655+1maple1@users.noreply.github.com>
Date: Mon, 19 Aug 2024 15:32:25 +0000
Subject: [PATCH] Add requires timeout

---
 brewtils/models.py        | 2 ++
 brewtils/plugin.py        | 6 +++++-
 brewtils/schemas.py       | 1 +
 brewtils/specification.py | 5 +++++
 test/plugin_test.py       | 2 ++
 5 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/brewtils/models.py b/brewtils/models.py
index 3a067f44..d6578eea 100644
--- a/brewtils/models.py
+++ b/brewtils/models.py
@@ -763,6 +763,7 @@ def __init__(
         groups=None,
         prefix_topic=None,
         requires=None,
+        requires_timeout=None,
     ):
         self.name = name
         self.description = description
@@ -780,6 +781,7 @@ def __init__(
         self.groups = groups or []
         self.prefix_topic = prefix_topic
         self.requires = requires or []
+        self.requires_timeout = requires_timeout
 
     def __str__(self):
         return "%s:%s-%s" % (self.namespace, self.name, self.version)
diff --git a/brewtils/plugin.py b/brewtils/plugin.py
index 4813040e..526c3ca2 100644
--- a/brewtils/plugin.py
+++ b/brewtils/plugin.py
@@ -87,6 +87,7 @@ class Plugin(object):
         - ``groups``
         - ``require``
         - ``requires``
+        - ``requires_timeout``
 
     Connection information tells the Plugin how to communicate with Beer-garden. The
     most important of these is the ``bg_host`` (to tell the plugin where to find the
@@ -176,6 +177,7 @@ class Plugin(object):
 
         require (str): Required system dependency
         requires (list): Required systems dependencies
+        requires_timeout (int): Timeout to wait for dependencies
 
         group (str): Grouping label applied to plugin
         groups (list): Grouping labels applied to plugin
@@ -401,7 +403,7 @@ def get_system_dependency(self, require, timeout=300):
 
     def await_dependencies(self, config):
         for req in config.requires:
-            system = self.get_system_dependency(req)
+            system = self.get_system_dependency(req, config.requires_timeout)
             self.logger.info(
                 f"Resolved system {system} for {req}: {config.name} {config.instance_name}"
             )
@@ -446,6 +448,7 @@ def _startup(self):
         self._admin_processor.startup()
 
         try:
+            print(self._config)
             if self._config.requires:
                 self.await_dependencies(self._config)
         except PluginValidationError:
@@ -922,6 +925,7 @@ def _setup_system(self, system, plugin_kwargs):
                 groups=self._config.groups,
                 prefix_topic=self._config.prefix_topic,
                 requires=self._config.requires,
+                requires_timeout=self._config.requires_timeout,
             )
 
         return system
diff --git a/brewtils/schemas.py b/brewtils/schemas.py
index dbad8da1..085ae2bb 100644
--- a/brewtils/schemas.py
+++ b/brewtils/schemas.py
@@ -258,6 +258,7 @@ class SystemSchema(BaseSchema):
     groups = fields.List(fields.Str(), allow_none=True)
     prefix_topic = fields.Str(allow_none=True)
     requires = fields.List(fields.Str(), allow_none=True)
+    requires_timeout = fields.Integer(allow_none=True)
 
 
 class SystemDomainIdentifierSchema(BaseSchema):
diff --git a/brewtils/specification.py b/brewtils/specification.py
index 32e17497..2ae991c2 100644
--- a/brewtils/specification.py
+++ b/brewtils/specification.py
@@ -183,6 +183,11 @@ def _is_json_dict(s):
         "required": False,
         "default": [],
     },
+    "requires_timeout": {
+        "type": "int",
+        "description": "The dependency timeout to use",
+        "default": 300,
+    },
 }
 
 _PLUGIN_SPEC = {
diff --git a/test/plugin_test.py b/test/plugin_test.py
index 2cb251a7..5def7eba 100644
--- a/test/plugin_test.py
+++ b/test/plugin_test.py
@@ -171,6 +171,7 @@ def test_kwargs(self, client, bg_system):
             prefix_topic="custom.topic",
             require="SystemA",
             requires=["SystemB"],
+            requires_timeout=200,
         )
 
         assert plugin._logger == logger
@@ -186,6 +187,7 @@ def test_kwargs(self, client, bg_system):
         assert "SystemA" == plugin._config.require
         assert "SystemB" in plugin._config.requires
         assert "SystemA" not in plugin._config.requires
+        assert plugin._config.requires_timeout == 200
 
     def test_env(self, client, bg_system):
         os.environ["BG_HOST"] = "remotehost"