1 files changed, 3 insertions, 11 deletions
diff --git a/tests/unittests/test_datasource/test_azure_helper.py b/tests/unittests/test_datasource/test_azure_helper.py
index 1134199b..65202ff0 100644
--- a/tests/unittests/test_datasource/test_azure_helper.py
+++ b/tests/unittests/test_datasource/test_azure_helper.py
@@ -1,17 +1,8 @@
 import os
 
 from cloudinit.sources.helpers import azure as azure_helper
-from ..helpers import TestCase
 
-try:
-    from unittest import mock
-except ImportError:
-    import mock
-
-try:
-    from contextlib import ExitStack
-except ImportError:
-    from contextlib2 import ExitStack
+from ..helpers import ExitStack, mock, TestCase
 
 
 GOAL_STATE_TEMPLATE = """\
@@ -70,7 +61,7 @@ class TestFindEndpoint(TestCase):
 
     def test_missing_special_azure_line(self):
         self.load_file.return_value = ''
-        self.assertRaises(Exception,
+        self.assertRaises(ValueError,
                           azure_helper.WALinuxAgentShim.find_endpoint)
 
     @staticmethod
@@ -287,6 +278,7 @@ class TestOpenSSLManager(TestCase):
         self.subp.side_effect = capture_directory
         manager = azure_helper.OpenSSLManager()
         self.assertEqual(manager.tmpdir, subp_directory['path'])
+        manager.clean_up()
 
     @mock.patch.object(azure_helper, 'cd', mock.MagicMock())
     @mock.patch.object(azure_helper.tempfile, 'mkdtemp', mock.MagicMock())