From cb9ecd4bce0294354584a5b7587a3406eda6c383 Mon Sep 17 00:00:00 2001
From: David Barroso <dbarrosop@dravetech.com>
Date: Thu, 10 Aug 2017 10:51:47 +0200
Subject: [PATCH] __enter__ was swallowing the exception and returning False
 (#299)

* __enter__ was swallowing the exception and returning False

* Unit test for context manager exceptions

* Test exception handling in __open__

* Refactoring context manager
---
 napalm_base/base.py           | 40 ++++++++---------------------------
 napalm_base/mock.py           |  3 +++
 test/unit/test_mock_driver.py | 14 +++++++++++-
 3 files changed, 25 insertions(+), 32 deletions(-)

diff --git a/napalm_base/base.py b/napalm_base/base.py
index fcd9da3d..d337f951 100644
--- a/napalm_base/base.py
+++ b/napalm_base/base.py
@@ -16,9 +16,6 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-# std libs
-import sys
-
 # local modules
 import napalm_base.exceptions
 import napalm_base.helpers
@@ -47,17 +44,19 @@ def __init__(self, hostname, username, password, timeout=60, optional_args=None)
         raise NotImplementedError
 
     def __enter__(self):
-        try:
-            self.open()
-        except Exception:  # noqa
-            exc_info = sys.exc_info()
-            return self.__raise_clean_exception(exc_info[0], exc_info[1], exc_info[2])
+        self.open()
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         self.close()
-        if exc_type is not None:
-            return self.__raise_clean_exception(exc_type, exc_value, exc_traceback)
+        if exc_type is not None and (
+                            exc_type.__name__ not in dir(napalm_base.exceptions) and
+                            exc_type.__name__ not in __builtins__.keys()):
+            epilog = ("NAPALM didn't catch this exception. Please, fill a bugfix on "
+                      "https://github.com/napalm-automation/napalm/issues\n"
+                      "Don't forget to include this traceback.")
+            print(epilog)
+            return False
 
     def __del__(self):
         """
@@ -71,27 +70,6 @@ def __del__(self):
         except Exception:
             pass
 
-    @staticmethod
-    def __raise_clean_exception(exc_type, exc_value, exc_traceback):
-        """
-        This method is going to check if the exception exc_type is part of the builtins exceptions
-        or part of the napalm exceptions. If it is not, it will print a message on the screen
-        giving instructions to fill a bug.
-
-        Finally it will raise the original exception.
-
-        :param exc_type: Exception class.
-        :param exc_value: Exception object.
-        :param exc_traceback: Traceback.
-        """
-        if (exc_type.__name__ not in dir(napalm_base.exceptions) and
-                exc_type.__name__ not in __builtins__.keys()):
-            epilog = ("NAPALM didn't catch this exception. Please, fill a bugfix on "
-                      "https://github.com/napalm-automation/napalm/issues\n"
-                      "Don't forget to include this traceback.")
-            print(epilog)
-        return False
-
     def open(self):
         """
         Opens a connection to the device.
diff --git a/napalm_base/mock.py b/napalm_base/mock.py
index 18f08a42..cd2dc23c 100644
--- a/napalm_base/mock.py
+++ b/napalm_base/mock.py
@@ -106,6 +106,7 @@ def __init__(self, hostname, username, password, timeout=60, optional_args=None)
         self.password = password
         self.path = optional_args["path"]
         self.profile = optional_args.get("profile", [])
+        self.fail_on_open = optional_args.get("fail_on_open", False)
 
         self.opened = False
         self.calls = {}
@@ -126,6 +127,8 @@ def _raise_if_closed(self):
             raise napalm_base.exceptions.ConnectionClosedException("connection closed")
 
     def open(self):
+        if self.fail_on_open:
+            raise napalm_base.exceptions.ConnectionException("You told me to do this")
         self.opened = True
 
     def close(self):
diff --git a/test/unit/test_mock_driver.py b/test/unit/test_mock_driver.py
index f20f3ff8..5475e8b9 100644
--- a/test/unit/test_mock_driver.py
+++ b/test/unit/test_mock_driver.py
@@ -25,6 +25,11 @@
     "path": os.path.join(BASE_PATH, "test_mock_driver"),
     "profile": ["eos"],
 }
+fail_args = {
+    "path": os.path.join(BASE_PATH, "test_mock_driver"),
+    "profile": ["eos"],
+    "fail_on_open": True,
+}
 
 
 class TestMockDriver(object):
@@ -43,9 +48,16 @@ def test_basic(self):
         assert "connection closed" in py23_compat.text_type(excinfo.value)
 
     def test_context_manager(self):
-        with driver("blah", "bleh", "blih", optional_args=optional_args) as d:
+        with pytest.raises(napalm_base.exceptions.ConnectionException) as e, \
+                driver("blah", "bleh", "blih", optional_args=fail_args) as d:
+            pass
+        assert "You told me to do this" in py23_compat.text_type(e.value)
+        with pytest.raises(AttributeError) as e, \
+                driver("blah", "bleh", "blih", optional_args=optional_args) as d:
             assert d.is_alive() == {u'is_alive': True}
+            d.__fake_call()
         assert d.is_alive() == {u'is_alive': False}
+        assert "object has no attribute" in py23_compat.text_type(e.value)
 
     def test_mocking_getters(self):
         d = driver("blah", "bleh", "blih", optional_args=optional_args)