From 3219e14367c1e942e34cb8ae8354197f48d7c04b Mon Sep 17 00:00:00 2001
From: Matt Patrick <hazmat345@gmail.com>
Date: Wed, 7 Mar 2018 20:33:43 +0000
Subject: [PATCH] Fixing bug with multi-instance remote plugins

---
 brewtils/plugin.py  | 25 ++++++++++++++++++++++---
 test/plugin_test.py | 29 ++++++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/brewtils/plugin.py b/brewtils/plugin.py
index cea4b5b9..c0b5ca87 100644
--- a/brewtils/plugin.py
+++ b/brewtils/plugin.py
@@ -12,7 +12,8 @@
 
 import brewtils
 from brewtils.errors import BrewmasterValidationError, RequestProcessingError, \
-    DiscardMessageException, RepublishRequestException, BrewmasterConnectionError
+    DiscardMessageException, RepublishRequestException, BrewmasterConnectionError, \
+    PluginValidationError
 from brewtils.models import Instance, Request, System
 from brewtils.request_consumer import RequestConsumer
 from brewtils.rest.easy_client import EasyClient
@@ -320,6 +321,19 @@ def _initialize(self):
             else:
                 new_commands = None
 
+            if not existing_system.has_instance(self.instance_name):
+                if len(existing_system.instances) < existing_system.max_instances:
+                    existing_system.instances.append(Instance(name=self.instance_name))
+                    self.bm_client.create_system(existing_system)
+                else:
+                    raise PluginValidationError('Unable to create plugin with instance name "%s": '
+                                                'System "%s[%s]" has an instance limit of %d and '
+                                                'existing instances %s' %
+                                                (self.instance_name, existing_system.name,
+                                                 existing_system.version,
+                                                 existing_system.max_instances,
+                                                 ', '.join(existing_system.instance_names)))
+
             # We always update in case the metadata has changed.
             self.system = self.bm_client.update_system(existing_system.id,
                                                        new_commands=new_commands,
@@ -330,8 +344,13 @@ def _initialize(self):
         else:
             self.system = self.bm_client.create_system(self.system)
 
-        instance_id = next(instance.id for instance in self.system.instances
-                           if instance.name == self.instance_name)
+        # Sanity check to make sure an instance with this name was registered
+        if self.system.has_instance(self.instance_name):
+            instance_id = self.system.get_instance(self.instance_name).id
+        else:
+            raise PluginValidationError('Unable to find registered instance with name "%s"' %
+                                        self.instance_name)
+
         self.instance = self.bm_client.initialize_instance(instance_id)
 
         self.logger.debug("Plugin %s is initialized", self.unique_name)
diff --git a/test/plugin_test.py b/test/plugin_test.py
index 8e164cd4..3e03ae87 100644
--- a/test/plugin_test.py
+++ b/test/plugin_test.py
@@ -7,7 +7,7 @@
 from requests import ConnectionError
 
 from brewtils.errors import BrewmasterValidationError, RequestProcessingError, \
-    DiscardMessageException, RepublishRequestException
+    DiscardMessageException, RepublishRequestException, PluginValidationError
 from brewtils.models import Instance, Request, System, Command
 from brewtils.plugin import PluginBase
 
@@ -382,8 +382,7 @@ def test_initialize_system_exists_different_commands(self):
         self.bm_client_mock.update_system.return_value = self.system
 
         existing_system = System(id='id', name='test_system', version='1.0.0',
-                                 instances=[self.instance],
-                                 metadata={'foo': 'bar'})
+                                 instances=[self.instance], metadata={'foo': 'bar'})
         self.bm_client_mock.find_unique_system.return_value = existing_system
 
         self.plugin._initialize()
@@ -400,6 +399,24 @@ def test_initialize_system_exists_different_commands(self):
         self.assertEqual(self.plugin.system, self.bm_client_mock.create_system.return_value)
         self.assertEqual(self.plugin.instance, self.bm_client_mock.initialize_instance.return_value)
 
+    def test_initialize_system_new_instance(self):
+        self.plugin.instance_name = 'new_instance'
+
+        existing_system = System(id='id', name='test_system', version='1.0.0',
+                                 instances=[self.instance], max_instances=2,
+                                 metadata={'foo': 'bar'})
+        self.bm_client_mock.find_unique_system.return_value = existing_system
+
+        self.plugin._initialize()
+        self.assertTrue(self.bm_client_mock.create_system.called)
+        self.assertTrue(self.bm_client_mock.update_system.called)
+
+    def test_initialize_system_new_instance_maximum(self):
+        self.plugin.instance_name = 'new_instance'
+        self.bm_client_mock.find_unique_system.return_value = self.system
+
+        self.assertRaises(PluginValidationError, self.plugin._initialize)
+
     def test_initialize_system_update_metadata(self):
         self.system.commands = [Command('test')]
         self.bm_client_mock.update_system.return_value = self.system
@@ -421,6 +438,12 @@ def test_initialize_system_update_metadata(self):
         self.assertEqual(self.plugin.system, self.bm_client_mock.create_system.return_value)
         self.assertEqual(self.plugin.instance, self.bm_client_mock.initialize_instance.return_value)
 
+    def test_initialize_unregistered_instance(self):
+        self.system.has_instance = Mock(return_value=False)
+        self.bm_client_mock.find_unique_system.return_value = None
+
+        self.assertRaises(PluginValidationError, self.plugin._initialize)
+
     def test_shutdown(self):
         self.plugin.request_consumer = Mock()
         self.plugin.admin_consumer = Mock()