From 01457cedf6ce881b013cbedae07dbd95ecadb948 Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Thu, 1 Aug 2013 17:50:43 -0700
Subject: [PATCH 1/6] made service reconfiguration better, updated unit tests

---
 tests/core/service_test.py   | 76 ++++++++++++++++++++++++++++++++----
 tests/trond_test.py          |  2 +-
 tron/core/service.py         | 25 +++++++++---
 tron/core/serviceinstance.py | 19 +++++++++
 4 files changed, 109 insertions(+), 13 deletions(-)

diff --git a/tests/core/service_test.py b/tests/core/service_test.py
index 847ff7113..37f3501ef 100644
--- a/tests/core/service_test.py
+++ b/tests/core/service_test.py
@@ -149,6 +149,20 @@ def test_restore_state(self):
             self.instances.restore_state.return_value)
         self.service.enable.assert_called_with()
 
+    def test_update_node_pool_enabled(self):
+        autospec_method(self.service.repair)
+        self.service.enabled = True
+        self.service.update_node_pool()
+        self.service.instances.update_node_pool.assert_called_once_with()
+        self.service.repair.assert_called_once_with()
+
+    def test_update_node_pool_disabled(self):
+        autospec_method(self.service.repair)
+        self.service.enabled = False
+        self.service.update_node_pool()
+        self.service.instances.update_node_pool.assert_called_once_with()
+        assert not self.service.repair.called
+
 
 class ServiceCollectionTestCase(TestCase):
 
@@ -163,17 +177,65 @@ def _add_service(self):
             (serv.name, serv) for serv in self.service_list)
 
     @mock.patch('tron.core.service.Service', autospec=True)
-    def test_load_from_config(self, mock_service):
+    def test_build_with_new_config(self, service_patch):
+        new_config = mock.Mock(
+            name='i_come_from_the_land_of_ice_and_snow',
+            count=42)
+        old_service = mock.Mock(config=mock.Mock())
+        context = mock.create_autospec(command_context.CommandContext)
+        with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
+        as get_patch:
+            assert_equal(self.collection._build(new_config, context), service_patch.from_config('', ''))
+            service_patch.from_config.assert_any_call(new_config, context)
+            assert_equal(service_patch.from_config.call_count, 2)
+            get_patch.assert_called_once_with(new_config.name)
+
+    def test_build_with_same_config(self):
+        config = mock.Mock(
+            name='hamsteak',
+            count=413)
+        old_service = mock.Mock(config=config)
+        context = mock.create_autospec(command_context.CommandContext)
+        with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
+        as get_patch:
+            assert_equal(self.collection._build(config, context), old_service)
+            get_patch.assert_called_once_with(config.name)
+            old_service.update_node_pool.assert_called_once_with()
+
+    def test_build_with_diff_count(self):
+        name = 'ni'
+        old_config = mock.Mock(
+            count=77)
+        new_config = mock.Mock(
+            count=1111111111111)
+        new_eq = lambda s, o: (s.name == o.name and s.count == o.count)
+        old_config.__eq__ = new_eq
+        new_config.__eq__ = new_eq
+        # We have to do this, since name is an actual kwarg for mock.Mock().
+        old_config.name = name
+        new_config.name = name
+        old_service = mock.Mock(config=old_config)
+        context = mock.create_autospec(command_context.CommandContext)
+        with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
+        as get_patch:
+            assert_equal(self.collection._build(new_config, context), old_service)
+            get_patch.assert_called_once_with(name)
+            old_service.update_node_pool.assert_called_once_with()
+
+
+    def test_load_from_config(self):
         autospec_method(self.collection.get_names)
         autospec_method(self.collection.add)
+        autospec_method(self.collection._build)
         service_configs = {'a': mock.Mock(), 'b': mock.Mock()}
         context = mock.create_autospec(command_context.CommandContext)
-        result = list(self.collection.load_from_config(service_configs, context))
-        expected = [mock.call(config, context)
-                    for config in service_configs.itervalues()]
-        assert_mock_calls(expected, mock_service.from_config.mock_calls)
-        expected = [mock.call(s) for s in result]
-        assert_mock_calls(expected, self.collection.add.mock_calls)
+        with mock.patch.object(self.collection, '_build') as build_patch:
+            result = list(self.collection.load_from_config(service_configs, context))
+            expected = [mock.call(config, context)
+                        for config in service_configs.itervalues()]
+            build_patch.assert_calls(expected)
+            expected = [mock.call(s) for s in result]
+            assert_mock_calls(expected, self.collection.add.mock_calls)
 
     def test_add(self):
         self.collection.services = mock.MagicMock()
diff --git a/tests/trond_test.py b/tests/trond_test.py
index 57cdff306..cf26edde4 100644
--- a/tests/trond_test.py
+++ b/tests/trond_test.py
@@ -153,7 +153,7 @@ def test_node_reconfig(self):
         self.sandbox.tronfig(second_config)
 
         sandbox.wait_on_state(self.client.service, service_url,
-            service.ServiceState.DISABLED)
+            service.ServiceState.STARTING)
 
         job_url = self.client.get_url('MASTER.a_job')
         def wait_on_next_run():
diff --git a/tron/core/service.py b/tron/core/service.py
index c9493f3b7..ff678654c 100644
--- a/tron/core/service.py
+++ b/tron/core/service.py
@@ -156,6 +156,11 @@ def restore_state(self, state_data):
         (self.enable if state_data.get('enabled') else self.disable)()
         self.event_recorder.info("restored")
 
+    def update_node_pool(self):
+        self.instances.update_node_pool()
+        if self.enabled:
+            self.repair()
+
 
 class ServiceCollection(object):
     """A collection of services."""
@@ -163,17 +168,27 @@ class ServiceCollection(object):
     def __init__(self):
         self.services = collections.MappingCollection('services')
 
+    def _build(self, config, context):
+        old_service = self.get_by_name(config.name)
+        if old_service and old_service.config.count != config.count:
+            old_service.config.count = config.count
+
+        if old_service and old_service.config == config:
+            log.debug("Updating service %s\'s node pool" % config.name)
+            old_service.update_node_pool()
+            return old_service
+        else:
+            log.debug("Building new service %s", config.name)
+            return Service.from_config(config, context)
+
     def load_from_config(self, service_configs, context):
         """Apply a configuration to this collection and return a generator of
         services which were added.
         """
         self.services.filter_by_name(service_configs.keys())
 
-        def build(config):
-            log.debug("Building new service %s", config.name)
-            return Service.from_config(config, context)
-
-        seq = (build(config) for config in service_configs.itervalues())
+        seq = (self._build(config, context)
+            for config in service_configs.itervalues())
         return itertools.ifilter(self.add, seq)
 
     def add(self, service):
diff --git a/tron/core/serviceinstance.py b/tron/core/serviceinstance.py
index d56e0218f..ce042c5bd 100644
--- a/tron/core/serviceinstance.py
+++ b/tron/core/serviceinstance.py
@@ -461,6 +461,25 @@ def get_by_number(self, instance_number):
             if instance.instance_number == instance_number:
                 return instance
 
+    def update_node_pool(self):
+        """Attempt to load a new node pool from the NodePoolRepository, and
+        remove instances that no longer have their node in the NodePool."""
+        node_repo = node.NodePoolRepository.get_instance()
+        node_pool = node_repo.get_by_name(self.config.node)
+
+        if node_pool != self.node_pool:
+            self.node_pool = node_pool
+            needs_new_node = []
+
+            for instance in self.instances:
+                old_node = self.node_pool.get_by_name(instance.node)
+                if old_node != instance.node:
+                    instance.stop()
+                    needs_new_node.append(instance)
+
+            self.instances = [i for i in self.instances if
+                              i not in needs_new_node]
+
     @property
     def missing(self):
         return self.config.count - len(self.instances)

From b9e88c30b9648483010e76e224ec7c0d9047ee66 Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Fri, 2 Aug 2013 09:57:41 -0700
Subject: [PATCH 2/6] slight fix to update_node_pool, unit tests for
 update_node_pool

---
 tests/core/serviceinstance_test.py | 48 ++++++++++++++++++++++++++++++
 tron/core/serviceinstance.py       |  4 +--
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/tests/core/serviceinstance_test.py b/tests/core/serviceinstance_test.py
index ff75002e3..8b613805a 100644
--- a/tests/core/serviceinstance_test.py
+++ b/tests/core/serviceinstance_test.py
@@ -502,6 +502,54 @@ def test_get_by_number(self):
         instance = self.collection.get_by_number(3)
         assert_equal(instance, instances[3])
 
+    def test_update_node_pool_same_pool(self):
+        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
+            get_mock = pool_patch.get_instance().get_by_name
+            get_mock.configure_mock(return_value=self.collection.node_pool)
+            self.collection.update_node_pool()
+            assert_equal(pool_patch.get_instance.call_count, 2)
+            get_mock.assert_called_once_with(self.collection.config.node)
+
+    def test_update_node_pool_diff_pool_same_nodes(self):
+        new_instances = [mock.Mock(), mock.Mock()]
+        self.collection.instances = new_instances
+        nodes = [instance.node for instance in new_instances]
+        node_pool = mock.Mock(get_by_name=mock.Mock(side_effect=iter(nodes)))
+
+        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
+            get_mock = pool_patch.get_instance().get_by_name
+            get_mock.configure_mock(return_value=node_pool)
+
+            self.collection.update_node_pool()
+
+            assert_equal(pool_patch.get_instance.call_count, 2)
+            get_mock.assert_called_once_with(self.collection.config.node)
+            assert_equal(self.collection.node_pool, node_pool)
+            calls = [mock.call(instance.node.name) for instance in new_instances]
+            node_pool.get_by_name.assert_calls(calls)
+            assert not any([instance.stop.called for instance in new_instances])
+            assert_equal(self.collection.instances, new_instances)
+
+    def test_update_node_pool_diff_everything(self):
+        new_instances = [mock.Mock(), mock.Mock()]
+        self.collection.instances = [mock.Mock(), mock.Mock()]
+        nodes = [instance.node for instance in new_instances]
+        node_pool = mock.Mock(get_by_name=mock.Mock(side_effect=iter(nodes)))
+
+        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
+            get_mock = pool_patch.get_instance().get_by_name
+            get_mock.configure_mock(return_value=node_pool)
+
+            self.collection.update_node_pool()
+
+            assert_equal(pool_patch.get_instance.call_count, 2)
+            get_mock.assert_called_once_with(self.collection.config.node)
+            assert_equal(self.collection.node_pool, node_pool)
+            calls = [mock.call(instance.node.name) for instance in self.collection.instances]
+            node_pool.get_by_name.assert_calls(calls)
+            assert all([instance.stop.called for instance in self.collection.instances])
+            assert_equal(self.collection.instances, [])
+
 
 if __name__ == "__main__":
     run()
diff --git a/tron/core/serviceinstance.py b/tron/core/serviceinstance.py
index ce042c5bd..c6468dfa3 100644
--- a/tron/core/serviceinstance.py
+++ b/tron/core/serviceinstance.py
@@ -472,8 +472,8 @@ def update_node_pool(self):
             needs_new_node = []
 
             for instance in self.instances:
-                old_node = self.node_pool.get_by_name(instance.node)
-                if old_node != instance.node:
+                new_node = self.node_pool.get_by_name(instance.node.name)
+                if new_node != instance.node:
                     instance.stop()
                     needs_new_node.append(instance)
 

From 8b9bc0ad79a40a1f75a465490722dc69ca260a02 Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Fri, 2 Aug 2013 10:17:25 -0700
Subject: [PATCH 3/6] added clear_extra for removing instances when
 config.count decreases

---
 tests/core/service_test.py         |  2 ++
 tests/core/serviceinstance_test.py | 10 ++++++++++
 tron/core/service.py               |  1 +
 tron/core/serviceinstance.py       |  6 ++++++
 4 files changed, 19 insertions(+)

diff --git a/tests/core/service_test.py b/tests/core/service_test.py
index 37f3501ef..24d2a4afa 100644
--- a/tests/core/service_test.py
+++ b/tests/core/service_test.py
@@ -154,6 +154,7 @@ def test_update_node_pool_enabled(self):
         self.service.enabled = True
         self.service.update_node_pool()
         self.service.instances.update_node_pool.assert_called_once_with()
+        self.service.instances.clear_extra.assert_called_once_with()
         self.service.repair.assert_called_once_with()
 
     def test_update_node_pool_disabled(self):
@@ -161,6 +162,7 @@ def test_update_node_pool_disabled(self):
         self.service.enabled = False
         self.service.update_node_pool()
         self.service.instances.update_node_pool.assert_called_once_with()
+        self.service.instances.clear_extra.assert_called_once_with()
         assert not self.service.repair.called
 
 
diff --git a/tests/core/serviceinstance_test.py b/tests/core/serviceinstance_test.py
index 8b613805a..72ed3db1e 100644
--- a/tests/core/serviceinstance_test.py
+++ b/tests/core/serviceinstance_test.py
@@ -550,6 +550,16 @@ def test_update_node_pool_diff_everything(self):
             assert all([instance.stop.called for instance in self.collection.instances])
             assert_equal(self.collection.instances, [])
 
+    def test_clear_extra(self):
+        instance_a = mock.Mock()
+        instance_b = mock.Mock()
+        instance_c = mock.Mock()
+        self.collection.instances = [instance_a, instance_b, instance_c]
+        self.collection.config.count = 2
+        self.collection.clear_extra()
+        assert_equal(self.collection.instances, [instance_a, instance_b])
+        instance_c.stop.assert_called_once_with()
+
 
 if __name__ == "__main__":
     run()
diff --git a/tron/core/service.py b/tron/core/service.py
index ff678654c..58c010266 100644
--- a/tron/core/service.py
+++ b/tron/core/service.py
@@ -158,6 +158,7 @@ def restore_state(self, state_data):
 
     def update_node_pool(self):
         self.instances.update_node_pool()
+        self.instances.clear_extra()
         if self.enabled:
             self.repair()
 
diff --git a/tron/core/serviceinstance.py b/tron/core/serviceinstance.py
index c6468dfa3..e260973c5 100644
--- a/tron/core/serviceinstance.py
+++ b/tron/core/serviceinstance.py
@@ -480,6 +480,12 @@ def update_node_pool(self):
             self.instances = [i for i in self.instances if
                               i not in needs_new_node]
 
+    def clear_extra(self):
+        """Clear out instances if too many exist."""
+        for i in range(0, self.missing, -1):
+            instance = self.instances.pop()
+            instance.stop()
+
     @property
     def missing(self):
         return self.config.count - len(self.instances)

From 5cd9f1710cb6d836cf5dca0a8cbfdd861164c83f Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Fri, 2 Aug 2013 14:13:18 -0700
Subject: [PATCH 4/6] removed ServiceCollection.add, moved .get_instance call
 to Service instead of ServiceInstanceCollection

---
 tests/core/service_test.py         | 45 +++++++++++++++++++-----------
 tests/core/serviceinstance_test.py | 44 ++++++++++-------------------
 tron/core/service.py               | 10 +++----
 tron/core/serviceinstance.py       |  5 +---
 4 files changed, 48 insertions(+), 56 deletions(-)

diff --git a/tests/core/service_test.py b/tests/core/service_test.py
index 24d2a4afa..5868e190f 100644
--- a/tests/core/service_test.py
+++ b/tests/core/service_test.py
@@ -152,18 +152,36 @@ def test_restore_state(self):
     def test_update_node_pool_enabled(self):
         autospec_method(self.service.repair)
         self.service.enabled = True
-        self.service.update_node_pool()
-        self.service.instances.update_node_pool.assert_called_once_with()
-        self.service.instances.clear_extra.assert_called_once_with()
-        self.service.repair.assert_called_once_with()
+        node_pool = mock.Mock()
+
+        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
+            get_mock = pool_patch.get_instance().get_by_name
+            get_mock.configure_mock(return_value=node_pool)
+
+            self.service.update_node_pool()
+
+            assert_equal(pool_patch.get_instance.call_count, 2)
+            get_mock.assert_called_once_with(self.service.config.node)
+            self.service.instances.update_node_pool.assert_called_once_with(node_pool)
+            self.service.instances.clear_extra.assert_called_once_with()
+            self.service.repair.assert_called_once_with()
 
     def test_update_node_pool_disabled(self):
         autospec_method(self.service.repair)
         self.service.enabled = False
-        self.service.update_node_pool()
-        self.service.instances.update_node_pool.assert_called_once_with()
-        self.service.instances.clear_extra.assert_called_once_with()
-        assert not self.service.repair.called
+        node_pool = mock.Mock()
+
+        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
+            get_mock = pool_patch.get_instance().get_by_name
+            get_mock.configure_mock(return_value=node_pool)
+
+            self.service.update_node_pool()
+
+            assert_equal(pool_patch.get_instance.call_count, 2)
+            get_mock.assert_called_once_with(self.service.config.node)
+            self.service.instances.update_node_pool.assert_called_once_with(node_pool)
+            self.service.instances.clear_extra.assert_called_once_with()
+            assert not self.service.repair.called
 
 
 class ServiceCollectionTestCase(TestCase):
@@ -227,7 +245,7 @@ def test_build_with_diff_count(self):
 
     def test_load_from_config(self):
         autospec_method(self.collection.get_names)
-        autospec_method(self.collection.add)
+        autospec_method(self.collection.services.replace)
         autospec_method(self.collection._build)
         service_configs = {'a': mock.Mock(), 'b': mock.Mock()}
         context = mock.create_autospec(command_context.CommandContext)
@@ -237,14 +255,7 @@ def test_load_from_config(self):
                         for config in service_configs.itervalues()]
             build_patch.assert_calls(expected)
             expected = [mock.call(s) for s in result]
-            assert_mock_calls(expected, self.collection.add.mock_calls)
-
-    def test_add(self):
-        self.collection.services = mock.MagicMock()
-        service = mock.Mock()
-        result = self.collection.add(service)
-        self.collection.services.replace.assert_called_with(service)
-        assert_equal(result, self.collection.services.replace.return_value)
+            assert_mock_calls(expected, self.collection.services.replace.mock_calls)
 
     def test_restore_state(self):
         state_count = 2
diff --git a/tests/core/serviceinstance_test.py b/tests/core/serviceinstance_test.py
index 72ed3db1e..e4846a557 100644
--- a/tests/core/serviceinstance_test.py
+++ b/tests/core/serviceinstance_test.py
@@ -503,12 +503,8 @@ def test_get_by_number(self):
         assert_equal(instance, instances[3])
 
     def test_update_node_pool_same_pool(self):
-        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
-            get_mock = pool_patch.get_instance().get_by_name
-            get_mock.configure_mock(return_value=self.collection.node_pool)
-            self.collection.update_node_pool()
-            assert_equal(pool_patch.get_instance.call_count, 2)
-            get_mock.assert_called_once_with(self.collection.config.node)
+        self.collection.update_node_pool(self.collection.node_pool)
+        assert not self.collection.node_pool.get_by_name.called
 
     def test_update_node_pool_diff_pool_same_nodes(self):
         new_instances = [mock.Mock(), mock.Mock()]
@@ -516,19 +512,13 @@ def test_update_node_pool_diff_pool_same_nodes(self):
         nodes = [instance.node for instance in new_instances]
         node_pool = mock.Mock(get_by_name=mock.Mock(side_effect=iter(nodes)))
 
-        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
-            get_mock = pool_patch.get_instance().get_by_name
-            get_mock.configure_mock(return_value=node_pool)
+        self.collection.update_node_pool(node_pool)
 
-            self.collection.update_node_pool()
-
-            assert_equal(pool_patch.get_instance.call_count, 2)
-            get_mock.assert_called_once_with(self.collection.config.node)
-            assert_equal(self.collection.node_pool, node_pool)
-            calls = [mock.call(instance.node.name) for instance in new_instances]
-            node_pool.get_by_name.assert_calls(calls)
-            assert not any([instance.stop.called for instance in new_instances])
-            assert_equal(self.collection.instances, new_instances)
+        assert_equal(self.collection.node_pool, node_pool)
+        calls = [mock.call(instance.node.name) for instance in new_instances]
+        node_pool.get_by_name.assert_calls(calls)
+        assert not any([instance.stop.called for instance in new_instances])
+        assert_equal(self.collection.instances, new_instances)
 
     def test_update_node_pool_diff_everything(self):
         new_instances = [mock.Mock(), mock.Mock()]
@@ -536,19 +526,13 @@ def test_update_node_pool_diff_everything(self):
         nodes = [instance.node for instance in new_instances]
         node_pool = mock.Mock(get_by_name=mock.Mock(side_effect=iter(nodes)))
 
-        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
-            get_mock = pool_patch.get_instance().get_by_name
-            get_mock.configure_mock(return_value=node_pool)
-
-            self.collection.update_node_pool()
+        self.collection.update_node_pool(node_pool)
 
-            assert_equal(pool_patch.get_instance.call_count, 2)
-            get_mock.assert_called_once_with(self.collection.config.node)
-            assert_equal(self.collection.node_pool, node_pool)
-            calls = [mock.call(instance.node.name) for instance in self.collection.instances]
-            node_pool.get_by_name.assert_calls(calls)
-            assert all([instance.stop.called for instance in self.collection.instances])
-            assert_equal(self.collection.instances, [])
+        assert_equal(self.collection.node_pool, node_pool)
+        calls = [mock.call(instance.node.name) for instance in self.collection.instances]
+        node_pool.get_by_name.assert_calls(calls)
+        assert all([instance.stop.called for instance in self.collection.instances])
+        assert_equal(self.collection.instances, [])
 
     def test_clear_extra(self):
         instance_a = mock.Mock()
diff --git a/tron/core/service.py b/tron/core/service.py
index 58c010266..94e37332b 100644
--- a/tron/core/service.py
+++ b/tron/core/service.py
@@ -157,7 +157,10 @@ def restore_state(self, state_data):
         self.event_recorder.info("restored")
 
     def update_node_pool(self):
-        self.instances.update_node_pool()
+        node_repo = node.NodePoolRepository.get_instance()
+        node_pool = node_repo.get_by_name(self.config.node)
+
+        self.instances.update_node_pool(node_pool)
         self.instances.clear_extra()
         if self.enabled:
             self.repair()
@@ -190,10 +193,7 @@ def load_from_config(self, service_configs, context):
 
         seq = (self._build(config, context)
             for config in service_configs.itervalues())
-        return itertools.ifilter(self.add, seq)
-
-    def add(self, service):
-        return self.services.replace(service)
+        return itertools.ifilter(self.services.replace, seq)
 
     def restore_state(self, service_state_data):
         self.services.restore_state(service_state_data)
diff --git a/tron/core/serviceinstance.py b/tron/core/serviceinstance.py
index e260973c5..5e063a397 100644
--- a/tron/core/serviceinstance.py
+++ b/tron/core/serviceinstance.py
@@ -461,12 +461,9 @@ def get_by_number(self, instance_number):
             if instance.instance_number == instance_number:
                 return instance
 
-    def update_node_pool(self):
+    def update_node_pool(self, node_pool):
         """Attempt to load a new node pool from the NodePoolRepository, and
         remove instances that no longer have their node in the NodePool."""
-        node_repo = node.NodePoolRepository.get_instance()
-        node_pool = node_repo.get_by_name(self.config.node)
-
         if node_pool != self.node_pool:
             self.node_pool = node_pool
             needs_new_node = []

From 52993edf352b5432bfce115911df0b547e9e84bf Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Fri, 2 Aug 2013 16:11:52 -0700
Subject: [PATCH 5/6] factored _build into an update_func for collections.add

---
 tests/core/service_test.py | 44 ++++++++++++++++++--------------------
 tests/trond_test.py        |  2 +-
 tron/core/service.py       | 26 +++++++++++-----------
 3 files changed, 36 insertions(+), 36 deletions(-)

diff --git a/tests/core/service_test.py b/tests/core/service_test.py
index 5868e190f..8a39bdbbb 100644
--- a/tests/core/service_test.py
+++ b/tests/core/service_test.py
@@ -196,18 +196,16 @@ def _add_service(self):
         self.collection.services.update(
             (serv.name, serv) for serv in self.service_list)
 
-    @mock.patch('tron.core.service.Service', autospec=True)
-    def test_build_with_new_config(self, service_patch):
+    def test_build_with_new_config(self):
         new_config = mock.Mock(
             name='i_come_from_the_land_of_ice_and_snow',
             count=42)
-        old_service = mock.Mock(config=mock.Mock())
-        context = mock.create_autospec(command_context.CommandContext)
+        new_service = mock.Mock(config=new_config)
+        old_service = mock.Mock()
         with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
         as get_patch:
-            assert_equal(self.collection._build(new_config, context), service_patch.from_config('', ''))
-            service_patch.from_config.assert_any_call(new_config, context)
-            assert_equal(service_patch.from_config.call_count, 2)
+            assert not self.collection._build(new_service)
+            assert not old_service.update_node_pool.called
             get_patch.assert_called_once_with(new_config.name)
 
     def test_build_with_same_config(self):
@@ -215,12 +213,13 @@ def test_build_with_same_config(self):
             name='hamsteak',
             count=413)
         old_service = mock.Mock(config=config)
-        context = mock.create_autospec(command_context.CommandContext)
+        new_service = mock.Mock(config=config)
         with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
         as get_patch:
-            assert_equal(self.collection._build(config, context), old_service)
+            assert self.collection._build(new_service)
             get_patch.assert_called_once_with(config.name)
             old_service.update_node_pool.assert_called_once_with()
+            assert_equal(old_service.instances.context, new_service.instances.context)
 
     def test_build_with_diff_count(self):
         name = 'ni'
@@ -235,27 +234,26 @@ def test_build_with_diff_count(self):
         old_config.name = name
         new_config.name = name
         old_service = mock.Mock(config=old_config)
-        context = mock.create_autospec(command_context.CommandContext)
+        new_service = mock.Mock(config=new_config)
         with mock.patch.object(self.collection, 'get_by_name', return_value=old_service) \
         as get_patch:
-            assert_equal(self.collection._build(new_config, context), old_service)
-            get_patch.assert_called_once_with(name)
+            assert self.collection._build(new_service)
+            get_patch.assert_called_once_with(new_service.config.name)
             old_service.update_node_pool.assert_called_once_with()
+            assert_equal(old_service.instances.context, new_service.instances.context)
 
-
-    def test_load_from_config(self):
+    @mock.patch('tron.core.service.Service', autospec=True)
+    def test_load_from_config(self, mock_service):
         autospec_method(self.collection.get_names)
-        autospec_method(self.collection.services.replace)
-        autospec_method(self.collection._build)
+        autospec_method(self.collection.services.add)
         service_configs = {'a': mock.Mock(), 'b': mock.Mock()}
         context = mock.create_autospec(command_context.CommandContext)
-        with mock.patch.object(self.collection, '_build') as build_patch:
-            result = list(self.collection.load_from_config(service_configs, context))
-            expected = [mock.call(config, context)
-                        for config in service_configs.itervalues()]
-            build_patch.assert_calls(expected)
-            expected = [mock.call(s) for s in result]
-            assert_mock_calls(expected, self.collection.services.replace.mock_calls)
+        result = list(self.collection.load_from_config(service_configs, context))
+        expected = [mock.call(config, context)
+                    for config in service_configs.itervalues()]
+        assert_mock_calls(expected, mock_service.from_config.mock_calls)
+        expected = [mock.call(s, self.collection._build) for s in result]
+        assert_mock_calls(expected, self.collection.services.add.mock_calls)
 
     def test_restore_state(self):
         state_count = 2
diff --git a/tests/trond_test.py b/tests/trond_test.py
index cf26edde4..3aae26b17 100644
--- a/tests/trond_test.py
+++ b/tests/trond_test.py
@@ -153,7 +153,7 @@ def test_node_reconfig(self):
         self.sandbox.tronfig(second_config)
 
         sandbox.wait_on_state(self.client.service, service_url,
-            service.ServiceState.STARTING)
+            service.ServiceState.FAILED)
 
         job_url = self.client.get_url('MASTER.a_job')
         def wait_on_next_run():
diff --git a/tron/core/service.py b/tron/core/service.py
index 94e37332b..1d7086e4a 100644
--- a/tron/core/service.py
+++ b/tron/core/service.py
@@ -172,18 +172,20 @@ class ServiceCollection(object):
     def __init__(self):
         self.services = collections.MappingCollection('services')
 
-    def _build(self, config, context):
-        old_service = self.get_by_name(config.name)
-        if old_service and old_service.config.count != config.count:
-            old_service.config.count = config.count
-
-        if old_service and old_service.config == config:
-            log.debug("Updating service %s\'s node pool" % config.name)
+    def _build(self, new_service):
+        old_service = self.get_by_name(new_service.config.name)
+        if old_service and old_service.config.count != new_service.config.count:
+            old_service.config.count = new_service.config.count
+
+        if old_service and old_service.config == new_service.config:
+            log.debug("Updating service %s\'s node pool" % new_service.config.name)
+            old_service.instances.context = new_service.instances.context
             old_service.update_node_pool()
-            return old_service
+            return True
         else:
-            log.debug("Building new service %s", config.name)
-            return Service.from_config(config, context)
+            log.debug("Building new service %s", new_service.config.name)
+            old_service.disable()
+            return False
 
     def load_from_config(self, service_configs, context):
         """Apply a configuration to this collection and return a generator of
@@ -191,9 +193,9 @@ def load_from_config(self, service_configs, context):
         """
         self.services.filter_by_name(service_configs.keys())
 
-        seq = (self._build(config, context)
+        seq = (Service.from_config(config, context)
             for config in service_configs.itervalues())
-        return itertools.ifilter(self.services.replace, seq)
+        return itertools.ifilter(lambda e: self.services.add(e, self._build), seq)
 
     def restore_state(self, service_state_data):
         self.services.restore_state(service_state_data)

From a991b80f93f95e71c76e2d98aa25b210b7bb7ba9 Mon Sep 17 00:00:00 2001
From: James McGuinness <jrm@yelp.com>
Date: Tue, 6 Aug 2013 17:13:47 -0700
Subject: [PATCH 6/6] improved _build, update_node_pool to be more pythonic

---
 tests/core/service_test.py   | 32 ++++++++++----------------------
 tron/core/service.py         | 27 ++++++++++++++++++++-------
 tron/core/serviceinstance.py | 14 ++++++++------
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/tests/core/service_test.py b/tests/core/service_test.py
index 8a39bdbbb..01f12371e 100644
--- a/tests/core/service_test.py
+++ b/tests/core/service_test.py
@@ -154,34 +154,22 @@ def test_update_node_pool_enabled(self):
         self.service.enabled = True
         node_pool = mock.Mock()
 
-        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
-            get_mock = pool_patch.get_instance().get_by_name
-            get_mock.configure_mock(return_value=node_pool)
+        self.service.update_node_pool(node_pool)
 
-            self.service.update_node_pool()
-
-            assert_equal(pool_patch.get_instance.call_count, 2)
-            get_mock.assert_called_once_with(self.service.config.node)
-            self.service.instances.update_node_pool.assert_called_once_with(node_pool)
-            self.service.instances.clear_extra.assert_called_once_with()
-            self.service.repair.assert_called_once_with()
+        self.service.instances.update_node_pool.assert_called_once_with(node_pool)
+        self.service.instances.clear_extra.assert_called_once_with()
+        self.service.repair.assert_called_once_with()
 
     def test_update_node_pool_disabled(self):
         autospec_method(self.service.repair)
         self.service.enabled = False
         node_pool = mock.Mock()
 
-        with mock.patch('tron.node.NodePoolRepository', autospec=True) as pool_patch:
-            get_mock = pool_patch.get_instance().get_by_name
-            get_mock.configure_mock(return_value=node_pool)
-
-            self.service.update_node_pool()
+        self.service.update_node_pool(node_pool)
 
-            assert_equal(pool_patch.get_instance.call_count, 2)
-            get_mock.assert_called_once_with(self.service.config.node)
-            self.service.instances.update_node_pool.assert_called_once_with(node_pool)
-            self.service.instances.clear_extra.assert_called_once_with()
-            assert not self.service.repair.called
+        self.service.instances.update_node_pool.assert_called_once_with(node_pool)
+        self.service.instances.clear_extra.assert_called_once_with()
+        assert not self.service.repair.called
 
 
 class ServiceCollectionTestCase(TestCase):
@@ -218,7 +206,7 @@ def test_build_with_same_config(self):
         as get_patch:
             assert self.collection._build(new_service)
             get_patch.assert_called_once_with(config.name)
-            old_service.update_node_pool.assert_called_once_with()
+            old_service.update_node_pool.assert_called_once_with(new_service.instances.node_pool)
             assert_equal(old_service.instances.context, new_service.instances.context)
 
     def test_build_with_diff_count(self):
@@ -239,7 +227,7 @@ def test_build_with_diff_count(self):
         as get_patch:
             assert self.collection._build(new_service)
             get_patch.assert_called_once_with(new_service.config.name)
-            old_service.update_node_pool.assert_called_once_with()
+            old_service.update_node_pool.assert_called_once_with(new_service.instances.node_pool)
             assert_equal(old_service.instances.context, new_service.instances.context)
 
     @mock.patch('tron.core.service.Service', autospec=True)
diff --git a/tron/core/service.py b/tron/core/service.py
index 1d7086e4a..1034a4c2f 100644
--- a/tron/core/service.py
+++ b/tron/core/service.py
@@ -156,10 +156,7 @@ def restore_state(self, state_data):
         (self.enable if state_data.get('enabled') else self.disable)()
         self.event_recorder.info("restored")
 
-    def update_node_pool(self):
-        node_repo = node.NodePoolRepository.get_instance()
-        node_pool = node_repo.get_by_name(self.config.node)
-
+    def update_node_pool(self, node_pool):
         self.instances.update_node_pool(node_pool)
         self.instances.clear_extra()
         if self.enabled:
@@ -173,14 +170,30 @@ def __init__(self):
         self.services = collections.MappingCollection('services')
 
     def _build(self, new_service):
+        """A method to be used as an update function for MappingCollection.add.
+        This function attempts to load an old Service object, and if one
+        exists, see if we don't actually have to use an entirely new
+        Service object on reconfiguration.
+
+        To do this, we first check if the number of instances (config.count) is
+        different, as we have a method to fix this when updating the service's
+        node pool. Then, if the configs are now equal, we can simply update
+        the node pool of the old Service object and be done- no need for the
+        new Service object. Otherwise, we use the new object as normal.
+        """
         old_service = self.get_by_name(new_service.config.name)
-        if old_service and old_service.config.count != new_service.config.count:
+
+        if not old_service:
+            log.debug("Building new service %s", new_service.config.name)
+            return False
+
+        if old_service.config.count != new_service.config.count:
             old_service.config.count = new_service.config.count
 
-        if old_service and old_service.config == new_service.config:
+        if old_service.config == new_service.config:
             log.debug("Updating service %s\'s node pool" % new_service.config.name)
             old_service.instances.context = new_service.instances.context
-            old_service.update_node_pool()
+            old_service.update_node_pool(new_service.instances.node_pool)
             return True
         else:
             log.debug("Building new service %s", new_service.config.name)
diff --git a/tron/core/serviceinstance.py b/tron/core/serviceinstance.py
index 5e063a397..abc34b8e6 100644
--- a/tron/core/serviceinstance.py
+++ b/tron/core/serviceinstance.py
@@ -464,18 +464,20 @@ def get_by_number(self, instance_number):
     def update_node_pool(self, node_pool):
         """Attempt to load a new node pool from the NodePoolRepository, and
         remove instances that no longer have their node in the NodePool."""
-        if node_pool != self.node_pool:
-            self.node_pool = node_pool
-            needs_new_node = []
+        if node_pool == self.node_pool:
+            return
+
+        self.node_pool = node_pool
 
+        def _trim_old_nodes():
             for instance in self.instances:
                 new_node = self.node_pool.get_by_name(instance.node.name)
                 if new_node != instance.node:
                     instance.stop()
-                    needs_new_node.append(instance)
+                else:
+                    yield instance
 
-            self.instances = [i for i in self.instances if
-                              i not in needs_new_node]
+        self.instances = list(_trim_old_nodes())
 
     def clear_extra(self):
         """Clear out instances if too many exist."""