fix broker select error in arbiter

shinken/daemons/arbiterdaemon.py

@@ -217,13 +217,24 @@ def add(self, b):
     # TODO: better find the broker, here it can be dead?
     # or not the good one?
     def push_broks_to_broker(self):
-        for brk in self.conf.brokers:
-            # Send only if alive of course
-            if brk.manage_arbiters and brk.alive:
-                is_send = brk.push_broks(self.broks)
-                if is_send:
-                    # They are gone, we keep none!
-                    self.broks.clear()
+        # we may have less master arbiters than master schedulers, in cases
+        # where we have many realms with many schedulers with only few brokers,
+        # which is computed as below:
+        #   scheduler_len = 0
+        #   for r in self.realms:
+        #       scheduler_len += len(r.confs)
+        # so we may have duplicated brokers if we do not filter.
+
+        brks = set()
+        for r in self.conf.realms:
+            for cfg_id in r.confs:
+                tmp_brks = r.to_satellites_managed_by['broker'][cfg_id]
+                for tmp_brk in tmp_brks:
+                    brks.add(tmp_brk)
+        for brk in brks:
+            brk.push_broks(self.broks)
+        # They are gone, we keep none!
+        self.broks.clear()
 
     # We must take external_commands from all satellites
     # like brokers, pollers, reactionners or receivers

shinken/misc/regenerator.py

@@ -933,50 +933,53 @@ def manage_update_broker_status_brok(self, b):
         broker_name = data['broker_name']
         try:
             s = self.brokers[broker_name]
+        except KeyError:
+            self.manage_initial_broker_status_brok(b)
+        else:
             self.update_element(s, data)
-        except Exception:
-            pass
 
 
     def manage_update_receiver_status_brok(self, b):
         data = b.data
         receiver_name = data['receiver_name']
         try:
             s = self.receivers[receiver_name]
+        except KeyError:
+            self.manage_initial_receiver_status_brok(b)
+        else:
             self.update_element(s, data)
-        except Exception:
-            pass
-
 
     def manage_update_reactionner_status_brok(self, b):
         data = b.data
         reactionner_name = data['reactionner_name']
         try:
             s = self.reactionners[reactionner_name]
+        except KeyError:
+            self.manage_initial_reactionner_status_brok(b)
+        else:
             self.update_element(s, data)
-        except Exception:
-            pass
 
 
     def manage_update_poller_status_brok(self, b):
         data = b.data
         poller_name = data['poller_name']
         try:
             s = self.pollers[poller_name]
+        except KeyError:
+            self.manage_initial_poller_status_brok(b)
+        else:
             self.update_element(s, data)
-        except Exception:
-            pass
 
 
     def manage_update_scheduler_status_brok(self, b):
         data = b.data
         scheduler_name = data['scheduler_name']
         try:
             s = self.schedulers[scheduler_name]
+        except KeyError:
+            self.manage_initial_scheduler_status_brok(b)
+        else:
             self.update_element(s, data)
-            # print "S:", s
-        except Exception:
-            pass
 
 
 #################

test/test_regenerator_with_broker_restart.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright (C) 2009-2014:
+#    Gabes Jean, [email protected]
+#    Gerhard Lausser, [email protected]
+#
+# This file is part of Shinken.
+#
+# Shinken is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Shinken is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with Shinken.  If not, see <http://www.gnu.org/licenses/>.
+
+#
+# This file is used to test reading and processing of config files
+#
+
+import mock
+
+from shinken_test import *
+from shinken.misc import regenerator
+from shinken import brok
+from shinken.objects import item
+
+class TestConfig(ShinkenTest):
+    # setUp is inherited from ShinkenTest
+    def _get_update_broker_brok(self):
+        b = brok.Brok('update_broker_status', {'broker_name': 'broker-master'})
+        b.prepare()
+        return b
+
+    def _get_update_scheduler_brok(self):
+        b = brok.Brok('update_scheduler_status', {'scheduler_name':
+                                                      'scheduler-master'})
+        b.prepare()
+        return b
+
+    def _get_update_poller_brok(self):
+        b = brok.Brok('update_poller_status', {'poller_name': 'poller-master'})
+        b.prepare()
+        return b
+
+    def _get_update_receiver_brok(self):
+        b = brok.Brok('update_receiver_status', {'receiver_name':
+                                                     'receiver-master'})
+        b.prepare()
+        return b
+
+    def _get_update_reactionner_brok(self):
+        b = brok.Brok('update_reactionner_status',
+                      {'reactionner_name':'reactionner-master'})
+        b.prepare()
+        return b
+
+    def test_manage_update_broker_status_brok_with_broker_restart(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_broker_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_broker_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None):
+                b = self._get_update_broker_brok()
+                reg = regenerator.Regenerator()
+                reg.manage_update_broker_status_brok(b)
+        mock_manage_initial_broker_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_broker_status_brok_with_broker_normally(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_broker_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_broker_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None) \
+                    as mock_update_element:
+                b = self._get_update_broker_brok()
+                with mock.patch.object(item.Items, 'index_item',
+                                       return_value=b.data):
+                    reg = regenerator.Regenerator()
+                    original_info = {}
+                    reg.brokers[b.data['broker_name']] = original_info
+                    reg.manage_update_broker_status_brok(b)
+        mock_update_element.assert_called_once_with(original_info, b.data)
+        with self.assertRaises(AssertionError):
+            mock_manage_initial_broker_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_scheduler_status_brok_with_broker_restart(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_scheduler_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_scheduler_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None):
+                b = self._get_update_scheduler_brok()
+                reg = regenerator.Regenerator()
+                reg.manage_update_scheduler_status_brok(b)
+        mock_manage_initial_scheduler_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_scheduler_status_brok_with_broker_normally(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_scheduler_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_scheduler_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None) \
+                    as mock_update_element:
+                b = self._get_update_scheduler_brok()
+                with mock.patch.object(item.Items, 'index_item',
+                                       return_value=b.data):
+                    reg = regenerator.Regenerator()
+                    original_info = {}
+                    reg.schedulers[b.data['scheduler_name']] = original_info
+                    reg.manage_update_scheduler_status_brok(b)
+        mock_update_element.assert_called_once_with(original_info, b.data)
+        with self.assertRaises(AssertionError):
+            mock_manage_initial_scheduler_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_poller_status_brok_with_broker_restart(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_poller_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_poller_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None):
+                b = self._get_update_poller_brok()
+                reg = regenerator.Regenerator()
+                reg.manage_update_poller_status_brok(b)
+        mock_manage_initial_poller_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_poller_status_brok_with_broker_normally(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_poller_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_poller_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None) \
+                    as mock_update_element:
+                b = self._get_update_poller_brok()
+                with mock.patch.object(item.Items, 'index_item',
+                                       return_value=b.data):
+                    reg = regenerator.Regenerator()
+                    original_info = {}
+                    reg.pollers[b.data['poller_name']] = original_info
+                    reg.manage_update_poller_status_brok(b)
+        mock_update_element.assert_called_once_with(original_info, b.data)
+        with self.assertRaises(AssertionError):
+            mock_manage_initial_poller_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_receiver_status_brok_with_broker_restart(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_receiver_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_receiver_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None):
+                b = self._get_update_receiver_brok()
+                reg = regenerator.Regenerator()
+                reg.manage_update_receiver_status_brok(b)
+        mock_manage_initial_receiver_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_receiver_status_brok_with_broker_normally(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_receiver_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_receiver_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None) \
+                    as mock_update_element:
+                b = self._get_update_receiver_brok()
+                with mock.patch.object(item.Items, 'index_item',
+                                       return_value=b.data):
+                    reg = regenerator.Regenerator()
+                    original_info = {}
+                    reg.receivers[b.data['receiver_name']] = original_info
+                    reg.manage_update_receiver_status_brok(b)
+        mock_update_element.assert_called_once_with(original_info, b.data)
+        with self.assertRaises(AssertionError):
+            mock_manage_initial_receiver_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_reactionner_status_brok_with_broker_restart(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_reactionner_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_reactionner_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None):
+                b = self._get_update_reactionner_brok()
+                reg = regenerator.Regenerator()
+                reg.manage_update_reactionner_status_brok(b)
+        mock_manage_initial_reactionner_status_brok.assert_called_once_with(b)
+
+    def test_manage_update_reactionner_status_brok_with_broker_normally(self):
+        with mock.patch.object(regenerator.Regenerator,
+                               'manage_initial_reactionner_status_brok',
+                               return_value=None) as \
+                mock_manage_initial_reactionner_status_brok:
+            with mock.patch.object(regenerator.Regenerator,
+                                   'update_element', return_value=None) \
+                    as mock_update_element:
+                b = self._get_update_reactionner_brok()
+                with mock.patch.object(item.Items, 'index_item',
+                                       return_value=b.data):
+                    reg = regenerator.Regenerator()
+                    original_info = {}
+                    reg.reactionners[b.data['reactionner_name']] = original_info
+                    reg.manage_update_reactionner_status_brok(b)
+        mock_update_element.assert_called_once_with(original_info, b.data)
+        with self.assertRaises(AssertionError):
+            mock_manage_initial_reactionner_status_brok\
+                .assert_called_once_with(b)
+
+if __name__ == '__main__':
+    unittest.main()