From 1ed7846b8da056ca8191fc10552d604df618b064 Mon Sep 17 00:00:00 2001
From: James Saryerwinnie <js@jamesls.com>
Date: Fri, 18 Sep 2015 15:34:13 -0700
Subject: [PATCH] Don't remove lazy component from deferred until success

Only delete the component from the deferred dict after
successfully creating the object from the factory as well as
injecting the instantiated value into the _components dict.

The bug is that is the factory object raises an exception,
we've already removed it from the deferred list (via the
.pop() call), but we haven't added anything to the _components
dict.

Because of this, the next time you try to make a call you'll
get an error message about an unknown component.

It also looks like there were no unit tests for this class.
It was only tested transitively before.  I've added a set
of unit tests covering this behavior.
---
 botocore/session.py        |  6 +++-
 tests/unit/test_session.py | 60 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/botocore/session.py b/botocore/session.py
index 411bdf62a8..ba88c040fa 100644
--- a/botocore/session.py
+++ b/botocore/session.py
@@ -770,8 +770,12 @@ def __init__(self):
 
     def get_component(self, name):
         if name in self._deferred:
-            factory = self._deferred.pop(name)
+            factory = self._deferred[name]
             self._components[name] = factory()
+            # Only delete the component from the deferred dict after
+            # successfully creating the object from the factory as well as
+            # injecting the instantiated value into the _components dict.
+            del self._deferred[name]
         try:
             return self._components[name]
         except KeyError:
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 4578929164..5bb58e18f1 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -449,5 +449,65 @@ def test_create_client_no_region_and_no_client_config(self):
         self.assertEqual(ec2_client.meta.region_name, 'moon-west-1')
 
 
+class TestComponentLocator(unittest.TestCase):
+    def setUp(self):
+        self.components = botocore.session.ComponentLocator()
+
+    def test_unknown_component_raises_exception(self):
+        with self.assertRaises(ValueError):
+            self.components.get_component('unknown-component')
+
+    def test_can_register_and_retrieve_component(self):
+        component = object()
+        self.components.register_component('foo', component)
+        self.assertIs(self.components.get_component('foo'), component)
+
+    def test_last_registration_wins(self):
+        first = object()
+        second = object()
+        self.components.register_component('foo', first)
+        self.components.register_component('foo', second)
+        self.assertIs(self.components.get_component('foo'), second)
+
+    def test_can_lazy_register_a_component(self):
+        component = object()
+        lazy = lambda: component
+        self.components.lazy_register_component('foo', lazy)
+        self.assertIs(self.components.get_component('foo'), component)
+
+    def test_latest_registration_wins_even_if_lazy(self):
+        first = object()
+        second = object()
+        lazy_second = lambda: second
+        self.components.register_component('foo', first)
+        self.components.lazy_register_component('foo', lazy_second)
+        self.assertIs(self.components.get_component('foo'), second)
+
+    def test_latest_registration_overrides_lazy(self):
+        first = object()
+        second = object()
+        lazy_first = lambda: first
+        self.components.lazy_register_component('foo', lazy_first)
+        self.components.register_component('foo', second)
+        self.assertIs(self.components.get_component('foo'), second)
+
+    def test_lazy_registration_factory_does_not_remove_from_list_on_error(self):
+        class ArbitraryError(Exception):
+            pass
+
+        def bad_factory():
+            raise ArbitraryError("Factory raises an exception.")
+
+        self.components.lazy_register_component('foo', bad_factory)
+
+        with self.assertRaises(ArbitraryError):
+            self.components.get_component('foo')
+
+        # Trying again should raise the same exception,
+        # not an ValueError("Unknown component")
+        with self.assertRaises(ArbitraryError):
+            self.components.get_component('foo')
+
+
 if __name__ == "__main__":
     unittest.main()