Repository: Azure/WALinuxAgent
Pull request #1704 (2019-11-07 18:41:57)

Remove 'from tests.tools import *'
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

The `tests.tools` module contains many helper functions used throughout other test modules. A lot of these test modules just used `from tests.tools import *`, which results in 1) importing unnecessary things, and 2) importing other agent modules implicitly via `tests.tools` instead of explicitly from the current test module.

<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] Except for special cases involving multiple contributors, the PR is started from a fork of the main repository, not a branch.
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made and Travis.CI is passing.

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/azure/walinuxagent/1704)
<!-- Reviewable:end -->


codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`develop@3d7ccb5`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704/graphs/tree.svg?width=650&token=ae90E2KYCw&height=150&src=pr)](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##             develop    #1704   +/-   ##
==========================================
  Coverage           ?   67.65%           
==========================================
  Files              ?       80           
  Lines              ?    11515           
  Branches           ?     1616           
==========================================
  Hits               ?     7790           
  Misses             ?     3383           
  Partials           ?      342
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704?src=pr&el=footer). Last update [3d7ccb5...3b1cc08](https://codecov.io/gh/Azure/WALinuxAgent/pull/1704?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 33f26e0b..ea6983f5 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -78,6 +78,7 @@ class CGroupConfigurator(object):
 
         def disable(self):
             self._enabled = False
+            CGroupsTelemetry.reset()
 
         def _invoke_cgroup_operation(self, operation, error_message, on_error=None):
             """
@@ -119,7 +120,6 @@ class CGroupConfigurator(object):
 
             def disable_cgroups(exception):
                 self.disable()
-                CGroupsTelemetry.reset()
                 add_event(
                     AGENT_NAME,
                     version=CURRENT_VERSION,
diff --git a/azurelinuxagent/common/dhcp.py b/azurelinuxagent/common/dhcp.py
index fc2cb649..59749656 100644
--- a/azurelinuxagent/common/dhcp.py
+++ b/azurelinuxagent/common/dhcp.py
@@ -19,8 +19,6 @@ import socket
 import array
 import time
 import azurelinuxagent.common.logger as logger
-import azurelinuxagent.common.utils.shellutil as shellutil
-from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.textutil import hex_dump, hex_dump2, \
     hex_dump3, \
     compare_bytes, str_to_ord, \
@@ -29,9 +27,10 @@ from azurelinuxagent.common.utils.textutil import hex_dump, hex_dump2, \
 from azurelinuxagent.common.exception import DhcpError
 from azurelinuxagent.common.osutil import get_osutil
 
+
 # the kernel routing table representation of 168.63.129.16
 KNOWN_WIRESERVER_IP_ENTRY = '10813FA8'
-KNOWN_WIRESERVER_IP = '168.63.129.16'
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
 
 
 def get_dhcp_handler():
@@ -152,10 +151,10 @@ class DhcpHandler(object):
         """
         Check if DHCP is available
         """
-        (dhcp_available, endpoint) =  self.osutil.is_dhcp_available()
+        dhcp_available =  self.osutil.is_dhcp_available()
         if not dhcp_available:
             logger.info("send_dhcp_req: DHCP not available")
-            self.endpoint = endpoint
+            self.endpoint = KNOWN_WIRESERVER_IP
             return
 
         """
diff --git a/azurelinuxagent/common/exception.py b/azurelinuxagent/common/exception.py
index 91c53fcb..b2cd30c5 100644
--- a/azurelinuxagent/common/exception.py
+++ b/azurelinuxagent/common/exception.py
@@ -52,14 +52,9 @@ class AgentNetworkError(AgentError):
 
 
 class CGroupsException(AgentError):
+    def __init__(self, msg=None, inner=None):
+        super(CGroupsException, self).__init__(msg, inner)
 
-    def __init__(self, msg, inner=None):
-        super(AgentError, self).__init__(msg, inner)
-        # TODO: AgentError should set the message - investigate whether doing it there would break anything
-        self.message = msg
-
-    def __str__(self):
-        return self.message
 
 class ExtensionError(AgentError):
     """
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index 7ccc5344..b99bca43 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -1072,7 +1072,7 @@ class DefaultOSUtil(object):
                       chk_err=False)
 
     def is_dhcp_available(self):
-        return (True, '')
+        return True
 
     def is_dhcp_enabled(self):
         return False
diff --git a/azurelinuxagent/common/osutil/iosxe.py b/azurelinuxagent/common/osutil/iosxe.py
index 3152ada8..0d45b0e2 100644
--- a/azurelinuxagent/common/osutil/iosxe.py
+++ b/azurelinuxagent/common/osutil/iosxe.py
@@ -63,7 +63,7 @@ class IosxeOSUtil(DefaultOSUtil):
         DefaultOSUtil.openssl_to_openssh(self, input_file, output_file)
 
     def is_dhcp_available(self):
-        return (False, '168.63.129.16')
+        return False
 
     def get_instance_id(self):
         '''
diff --git a/azurelinuxagent/common/protocol/util.py b/azurelinuxagent/common/protocol/util.py
index 692b5489..dab5a26a 100644
--- a/azurelinuxagent/common/protocol/util.py
+++ b/azurelinuxagent/common/protocol/util.py
@@ -35,9 +35,9 @@ from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.dhcp import get_dhcp_handler
 from azurelinuxagent.common.protocol.ovfenv import OvfEnv
 from azurelinuxagent.common.protocol.wire import WireProtocol
-from azurelinuxagent.common.protocol.metadata import MetadataProtocol, \
-                                                     METADATA_ENDPOINT
-from azurelinuxagent.common.utils.restutil import IOErrorCounter
+from azurelinuxagent.common.protocol.metadata import MetadataProtocol
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP, \
+                                                  IOErrorCounter
 
 OVF_FILE_NAME = "ovf-env.xml"
 TAG_FILE_NAME = "useMetadataEndpoint.tag"
@@ -64,14 +64,16 @@ def get_protocol_util():
 
 
 class ProtocolUtil(object):
-
     """
     ProtocolUtil handles initialization for protocol instance. 2 protocol types
     are invoked, wire protocol and metadata protocols.
     """
+
     def __init__(self):
         self.lock = threading.Lock()
+        self.lock_wireserver_endpoint = threading.Lock()
         self.protocol = None
+        self.endpoint = None
         self.osutil = get_osutil()
         self.dhcp_handler = get_dhcp_handler()
 
@@ -95,7 +97,7 @@ class ProtocolUtil(object):
         try:
             ovfxml = fileutil.read_file(ovf_file_path_on_dvd, remove_bom=True)
             ovfenv = OvfEnv(ovfxml)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error reading file "
                                 "{0}: {1}".format(ovf_file_path_on_dvd,
                                                   ustr(e)))
@@ -105,7 +107,7 @@ class ProtocolUtil(object):
                             PASSWORD_REPLACEMENT,
                             ovfxml)
             fileutil.write_file(ovf_file_path, ovfxml)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error writing file "
                                 "{0}: {1}".format(ovf_file_path,
                                                   ustr(e)))
@@ -114,7 +116,7 @@ class ProtocolUtil(object):
             if os.path.isfile(tag_file_path_on_dvd):
                 logger.info("Found {0} in provisioning ISO", TAG_FILE_NAME)
                 shutil.copyfile(tag_file_path_on_dvd, tag_file_path)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise ProtocolError("[CopyOvfEnv] Error copying file "
                                 "{0} to {1}: {2}".format(tag_file_path,
                                                          tag_file_path,
@@ -147,24 +149,77 @@ class ProtocolUtil(object):
             conf.get_lib_dir(),
             PROTOCOL_FILE_NAME)
 
+    def _get_wireserver_endpoint_file_path(self):
+        return os.path.join(
+            conf.get_lib_dir(),
+            ENDPOINT_FILE_NAME)
+
     def _get_tag_file_path(self):
         return os.path.join(
             conf.get_lib_dir(),
             TAG_FILE_NAME)
 
     def get_wireserver_endpoint(self):
+        self.lock_wireserver_endpoint.acquire()
+
         try:
-            file_path = os.path.join(conf.get_lib_dir(), ENDPOINT_FILE_NAME)
-            return fileutil.read_file(file_path)
-        except IOError as e:
-            raise OSUtilError(ustr(e))
+            if self.endpoint:
+                return self.endpoint
+
+            file_path = self._get_wireserver_endpoint_file_path()
+            if os.path.isfile(file_path):
+                try:
+                    self.endpoint = fileutil.read_file(file_path)
+
+                    if self.endpoint:
+                        logger.info("WireServer endpoint {0} read from file", self.endpoint)
+                        return self.endpoint
+
+                    logger.error("[GetWireserverEndpoint] Unexpected empty file {0}: {1}", file_path, str(e))
+                except (IOError, OSError) as e:
+                    logger.error("[GetWireserverEndpoint] Error reading file {0}: {1}", file_path, str(e))
+            else:
+                logger.error("[GetWireserverEndpoint] Missing file {0}", file_path)
+
+            self.endpoint = KNOWN_WIRESERVER_IP
+            logger.info("Using hardcoded Wireserver endpoint {0}", self.endpoint)
+
+            return self.endpoint
+        finally:
+            self.lock_wireserver_endpoint.release()
 
     def _set_wireserver_endpoint(self, endpoint):
+        self.lock_wireserver_endpoint.acquire()
         try:
-            file_path = os.path.join(conf.get_lib_dir(), ENDPOINT_FILE_NAME)
+            self.endpoint = endpoint
+            file_path = self._get_wireserver_endpoint_file_path()
             fileutil.write_file(file_path, endpoint)
-        except IOError as e:
+        except (IOError, OSError) as e:
             raise OSUtilError(ustr(e))
+        finally:
+            self.lock_wireserver_endpoint.release()
+
+    def _clear_wireserver_endpoint(self):
+        """
+        Cleanup previous saved wireserver endpoint.
+        """
+        self.lock_wireserver_endpoint.acquire()
+
+        try:
+            self.endpoint = None
+            endpoint_file_path = self._get_wireserver_endpoint_file_path()
+            if not os.path.isfile(endpoint_file_path):
+                return
+
+            try:
+                os.remove(endpoint_file_path)
+            except (IOError, OSError) as e:
+                # Ignore file-not-found errors (since the file is being removed)
+                if e.errno == errno.ENOENT:
+                    return
+                logger.error("Failed to clear wiresever endpoint: {0}", e)
+        finally:
+            self.lock_wireserver_endpoint.release()
 
     def _detect_wire_protocol(self):
         endpoint = self.dhcp_handler.endpoint
@@ -172,7 +227,7 @@ class ProtocolUtil(object):
             '''
             Check if DHCP can be used to get the wire protocol endpoint
             '''
-            (dhcp_available, conf_endpoint) =  self.osutil.is_dhcp_available()
+            dhcp_available = self.osutil.is_dhcp_available()
             if dhcp_available:
                 logger.info("WireServer endpoint is not found. Rerun dhcp handler")
                 try:
@@ -183,11 +238,6 @@ class ProtocolUtil(object):
             else:
                 logger.info("_detect_wire_protocol: DHCP not available")
                 endpoint = self.get_wireserver_endpoint()
-                if endpoint == None:
-                    endpoint = conf_endpoint
-                    logger.info("Using hardcoded WireServer endpoint {0}", endpoint)
-                else:
-                    logger.info("WireServer endpoint {0} read from file", endpoint)
 
         try:
             protocol = WireProtocol(endpoint)
@@ -195,7 +245,7 @@ class ProtocolUtil(object):
             self._set_wireserver_endpoint(endpoint)
             return protocol
         except ProtocolError as e:
-            logger.info("WireServer is not responding. Reset endpoint")
+            logger.info("WireServer is not responding. Reset dhcp endpoint")
             self.dhcp_handler.endpoint = None
             self.dhcp_handler.skip_cache = True
             raise e
@@ -254,14 +304,15 @@ class ProtocolUtil(object):
         protocol_file_path = self._get_protocol_file_path()
         try:
             fileutil.write_file(protocol_file_path, protocol_name)
-        except IOError as e:
+        except (IOError, OSError) as e:
             logger.error("Failed to save protocol endpoint: {0}", e)
 
     def clear_protocol(self):
         """
-        Cleanup previous saved endpoint.
+        Cleanup previous saved protocol endpoint.
         """
-        logger.info("Clean protocol")
+        logger.info("Clean protocol and wireserver endpoint")
+        self._clear_wireserver_endpoint()
         self.protocol = None
         protocol_file_path = self._get_protocol_file_path()
         if not os.path.isfile(protocol_file_path):
@@ -269,7 +320,7 @@ class ProtocolUtil(object):
 
         try:
             os.remove(protocol_file_path)
-        except IOError as e:
+        except (IOError, OSError) as e:
             # Ignore file-not-found errors (since the file is being removed)
             if e.errno == errno.ENOENT:
                 return
diff --git a/azurelinuxagent/common/utils/restutil.py b/azurelinuxagent/common/utils/restutil.py
index addba28b..fd78b969 100644
--- a/azurelinuxagent/common/utils/restutil.py
+++ b/azurelinuxagent/common/utils/restutil.py
@@ -98,13 +98,13 @@ HTTP_USER_AGENT_HEALTH = "{0}+health".format(HTTP_USER_AGENT)
 INVALID_CONTAINER_CONFIGURATION = "InvalidContainerConfiguration"
 REQUEST_ROLE_CONFIG_FILE_NOT_FOUND = "RequestRoleConfigFileNotFound"
 
-DEFAULT_PROTOCOL_ENDPOINT = '168.63.129.16'
+KNOWN_WIRESERVER_IP = '168.63.129.16'
 HOST_PLUGIN_PORT = 32526
 
 
 class IOErrorCounter(object):
     _lock = threading.RLock()
-    _protocol_endpoint = DEFAULT_PROTOCOL_ENDPOINT
+    _protocol_endpoint = KNOWN_WIRESERVER_IP
     _counts = {"hostplugin":0, "protocol":0, "other":0}
 
     @staticmethod
@@ -131,7 +131,7 @@ class IOErrorCounter(object):
             IOErrorCounter._counts = {"hostplugin":0, "protocol":0, "other":0}
 
     @staticmethod
-    def set_protocol_endpoint(endpoint=DEFAULT_PROTOCOL_ENDPOINT):
+    def set_protocol_endpoint(endpoint=KNOWN_WIRESERVER_IP):
         IOErrorCounter._protocol_endpoint = endpoint
 
 


Test patch:
diff --git a/tests/common/dhcp/test_dhcp.py b/tests/common/dhcp/test_dhcp.py
index 4be5e622..0b6adb0f 100644
--- a/tests/common/dhcp/test_dhcp.py
+++ b/tests/common/dhcp/test_dhcp.py
@@ -18,7 +18,7 @@
 import mock
 import azurelinuxagent.common.dhcp as dhcp
 import azurelinuxagent.common.osutil.default as osutil
-from tests.tools import *
+from tests.tools import AgentTestCase, open_patch, patch
 
 
 class TestDHCP(AgentTestCase):
diff --git a/tests/common/osutil/test_alpine.py b/tests/common/osutil/test_alpine.py
index d2f52b8d..ea0c6d8a 100644
--- a/tests/common/osutil/test_alpine.py
+++ b/tests/common/osutil/test_alpine.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.alpine import AlpineOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestAlpineOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_arch.py b/tests/common/osutil/test_arch.py
index f3c4b64a..eeda8560 100644
--- a/tests/common/osutil/test_arch.py
+++ b/tests/common/osutil/test_arch.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.arch import ArchUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestArchUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_bigip.py b/tests/common/osutil/test_bigip.py
index f8067eb6..43516167 100644
--- a/tests/common/osutil/test_bigip.py
+++ b/tests/common/osutil/test_bigip.py
@@ -15,16 +15,20 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import os
 import socket
+import time
+import unittest
 
 import azurelinuxagent.common.osutil.bigip as osutil
 import azurelinuxagent.common.osutil.default as default
 import azurelinuxagent.common.utils.shellutil as shellutil
 
 from azurelinuxagent.common.exception import OSUtilError
+import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.osutil.bigip import BigIpOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestBigIpOSUtil_wait_until_mcpd_is_initialized(AgentTestCase):
diff --git a/tests/common/osutil/test_clearlinux.py b/tests/common/osutil/test_clearlinux.py
index 03782f9b..a11d79a9 100644
--- a/tests/common/osutil/test_clearlinux.py
+++ b/tests/common/osutil/test_clearlinux.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.clearlinux import ClearLinuxUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestClearLinuxUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_coreos.py b/tests/common/osutil/test_coreos.py
index f5425873..fca73f6f 100644
--- a/tests/common/osutil/test_coreos.py
+++ b/tests/common/osutil/test_coreos.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.coreos import CoreOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestAlpineOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py
index 968499b7..f2dd87d1 100644
--- a/tests/common/osutil/test_default.py
+++ b/tests/common/osutil/test_default.py
@@ -19,16 +19,20 @@ import socket
 import glob
 import mock
 import traceback
-import re
+import os
+import tempfile
+import unittest
 
 import azurelinuxagent.common.osutil.default as osutil
 import azurelinuxagent.common.utils.shellutil as shellutil
 import azurelinuxagent.common.utils.textutil as textutil
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import OSUtilError
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil
-from tests.tools import *
-
+from azurelinuxagent.common.utils import fileutil
+from tests.tools import AgentTestCase, call, patch, open_patch, load_data, \
+    running_under_travis, skip_if_predicate_true
 
 actual_get_proc_net_route = 'azurelinuxagent.common.osutil.default.DefaultOSUtil._get_proc_net_route'
 
diff --git a/tests/common/osutil/test_default_osutil.py b/tests/common/osutil/test_default_osutil.py
index 2105e8c5..b0299ade 100644
--- a/tests/common/osutil/test_default_osutil.py
+++ b/tests/common/osutil/test_default_osutil.py
@@ -16,7 +16,8 @@
 #
 
 from azurelinuxagent.common.osutil.default import DefaultOSUtil, shellutil
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
+import os
 
 
 class DefaultOsUtilTestCase(AgentTestCase):
diff --git a/tests/common/osutil/test_factory.py b/tests/common/osutil/test_factory.py
index 92aa7d6e..adc099d5 100644
--- a/tests/common/osutil/test_factory.py
+++ b/tests/common/osutil/test_factory.py
@@ -32,7 +32,7 @@ from azurelinuxagent.common.osutil.bigip import BigIpOSUtil
 from azurelinuxagent.common.osutil.gaia import GaiaOSUtil
 from azurelinuxagent.common.osutil.iosxe import IosxeOSUtil
 from azurelinuxagent.common.osutil.openwrt import OpenWRTOSUtil
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestOsUtilFactory(AgentTestCase):
diff --git a/tests/common/osutil/test_freebsd.py b/tests/common/osutil/test_freebsd.py
index 1b1f687b..b7caa609 100644
--- a/tests/common/osutil/test_freebsd.py
+++ b/tests/common/osutil/test_freebsd.py
@@ -14,12 +14,14 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
-import mock
 
 from azurelinuxagent.common.osutil.freebsd import FreeBSDOSUtil
 import azurelinuxagent.common.utils.shellutil as shellutil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
+import traceback
+import unittest
+
 
 class TestFreeBSDOSUtil(AgentTestCase):
     def setUp(self):
diff --git a/tests/common/osutil/test_nsbsd.py b/tests/common/osutil/test_nsbsd.py
index 4e7082e5..87615f66 100644
--- a/tests/common/osutil/test_nsbsd.py
+++ b/tests/common/osutil/test_nsbsd.py
@@ -16,9 +16,9 @@
 #
 from azurelinuxagent.common.utils.fileutil import read_file
 from azurelinuxagent.common.osutil.nsbsd import NSBSDOSUtil
-from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 from os import path
+import unittest
 
 
 class TestNSBSDOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_openbsd.py b/tests/common/osutil/test_openbsd.py
index 76f214ef..1cee03f4 100644
--- a/tests/common/osutil/test_openbsd.py
+++ b/tests/common/osutil/test_openbsd.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.openbsd import OpenBSDOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestAlpineOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_openwrt.py b/tests/common/osutil/test_openwrt.py
index 3d72acd2..b6286dc9 100644
--- a/tests/common/osutil/test_openwrt.py
+++ b/tests/common/osutil/test_openwrt.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.openwrt import OpenWRTOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestOpenWRTOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_redhat.py b/tests/common/osutil/test_redhat.py
index ef8a3f96..c2926e03 100644
--- a/tests/common/osutil/test_redhat.py
+++ b/tests/common/osutil/test_redhat.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.redhat import Redhat6xOSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestRedhat6xOSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_suse.py b/tests/common/osutil/test_suse.py
index 8dbfce9c..e3ea661d 100644
--- a/tests/common/osutil/test_suse.py
+++ b/tests/common/osutil/test_suse.py
@@ -16,7 +16,8 @@
 #
 from azurelinuxagent.common.osutil.suse import SUSE11OSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
+import unittest
 
 
 class TestSUSE11OSUtil(AgentTestCase):
diff --git a/tests/common/osutil/test_ubuntu.py b/tests/common/osutil/test_ubuntu.py
index 6ca9aa2f..995c1c4b 100644
--- a/tests/common/osutil/test_ubuntu.py
+++ b/tests/common/osutil/test_ubuntu.py
@@ -14,9 +14,12 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
+
+import unittest
+
 from azurelinuxagent.common.osutil.ubuntu import Ubuntu12OSUtil, Ubuntu18OSUtil
 from .test_default import osutil_get_dhcp_pid_should_return_a_list_of_pids
-from tests.tools import *
+from tests.tools import AgentTestCase
 
 
 class TestUbuntu12OSUtil(AgentTestCase):
diff --git a/tests/common/test_cgroupapi.py b/tests/common/test_cgroupapi.py
index a88bd93f..bc3c0be9 100644
--- a/tests/common/test_cgroupapi.py
+++ b/tests/common/test_cgroupapi.py
@@ -637,7 +637,7 @@ class SystemdCgroupsApiMockedFileSystemTestCase(_MockedFileSystemTestCase):
                 with self.assertRaises(CGroupsException) as context_manager:
                     SystemdCgroupsApi().cleanup_legacy_cgroups()
 
-        self.assertEquals(context_manager.exception.message, "The daemon's PID ({0}) was already added to the legacy cgroup; this invalidates resource usage data.".format(daemon_pid))
+        self.assertEquals(str(context_manager.exception), "[CGroupsException] The daemon's PID ({0}) was already added to the legacy cgroup; this invalidates resource usage data.".format(daemon_pid))
 
         # The method should have deleted the legacy cgroups
         self.assertFalse(os.path.exists(legacy_cpu_cgroup))
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index af8182a1..fcc816cf 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -105,6 +105,18 @@ class CGroupConfiguratorTestCase(AgentTestCase):
                 CGroupConfigurator.get_instance().enable()
             self.assertIn("cgroups are not supported", str(context_manager.exception))
 
+    def test_disable_should_reset_tracked_cgroups(self):
+        configurator = CGroupConfigurator.get_instance()
+
+        # Start tracking a couple of dummy cgroups
+        CGroupsTelemetry.track_cgroup(CGroup("dummy", "/sys/fs/cgroup/memory/system.slice/dummy.service", "cpu"))
+        CGroupsTelemetry.track_cgroup(CGroup("dummy", "/sys/fs/cgroup/memory/system.slice/dummy.service", "memory"))
+
+        configurator.disable()
+
+        self.assertFalse(configurator.enabled())
+        self.assertEquals(len(CGroupsTelemetry._tracked), 0)
+
     def test_cgroup_operations_should_not_invoke_the_cgroup_api_when_cgroups_are_not_enabled(self):
         configurator = CGroupConfigurator.get_instance()
         configurator.disable()
@@ -291,7 +303,7 @@ class CGroupConfiguratorTestCase(AgentTestCase):
         self.assertFalse(kwargs['is_success'])
         self.assertEquals(
             kwargs['message'],
-            "Failed to process legacy cgroups. Collection of resource usage data will be disabled. The daemon's PID ({0}) was already added to the legacy cgroup; this invalidates resource usage data.".format(daemon_pid))
+            "Failed to process legacy cgroups. Collection of resource usage data will be disabled. [CGroupsException] The daemon's PID ({0}) was already added to the legacy cgroup; this invalidates resource usage data.".format(daemon_pid))
 
         self.assertFalse(cgroup_configurator.enabled())
         self.assertEquals(len(CGroupsTelemetry._tracked), 0)
diff --git a/tests/common/test_conf.py b/tests/common/test_conf.py
index b4805d34..725f7e57 100644
--- a/tests/common/test_conf.py
+++ b/tests/common/test_conf.py
@@ -15,12 +15,12 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-import mock
 import os.path
 
 from azurelinuxagent.common.conf import *
+import azurelinuxagent.common.conf as conf
 
-from tests.tools import *
+from tests.tools import AgentTestCase, data_dir, patch
 
 
 class TestConf(AgentTestCase):
diff --git a/tests/common/test_errorstate.py b/tests/common/test_errorstate.py
index a0a7761d..d578bdd9 100644
--- a/tests/common/test_errorstate.py
+++ b/tests/common/test_errorstate.py
@@ -1,7 +1,7 @@
-from datetime import timedelta
+import unittest
 
 from azurelinuxagent.common.errorstate import *
-from tests.tools import *
+from tests.tools import Mock, patch
 
 
 class TestErrorState(unittest.TestCase):
diff --git a/tests/common/test_event.py b/tests/common/test_event.py
index 4c676eee..9537b5ea 100644
--- a/tests/common/test_event.py
+++ b/tests/common/test_event.py
@@ -22,8 +22,6 @@ import os
 import threading
 from datetime import datetime, timedelta
 
-from mock import patch, Mock
-
 from azurelinuxagent.common import event, logger
 from azurelinuxagent.common.event import add_event, \
     WALAEventOperation, elapsed_milliseconds
@@ -34,7 +32,7 @@ from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.extensionprocessutil import read_output
 from azurelinuxagent.common.version import CURRENT_VERSION, CURRENT_AGENT
 from azurelinuxagent.ga.monitor import MonitorHandler
-from tests.tools import AgentTestCase, load_data, data_dir
+from tests.tools import AgentTestCase, load_data, data_dir, patch, Mock
 
 
 class TestEvent(AgentTestCase):
diff --git a/tests/common/test_exception.py b/tests/common/test_exception.py
new file mode 100644
index 00000000..2695d5fe
--- /dev/null
+++ b/tests/common/test_exception.py
@@ -0,0 +1,44 @@
+# Copyright 2019 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Requires Python 2.6+ and Openssl 1.0+
+#
+
+import sys, inspect
+from azurelinuxagent.common.exception import AgentError
+from tests.tools import AgentTestCase
+
+
+class TestAgentError(AgentTestCase):
+    @classmethod
+    def setUpClass(cls):
+        AgentTestCase.setUpClass()
+
+        cls.agent_exceptions = inspect.getmembers(
+            sys.modules["azurelinuxagent.common.exception"],
+            lambda member: inspect.isclass(member) and issubclass(member, AgentError))
+
+    def test_agent_exceptions_should_set_their_error_message(self):
+        for exception_name, exception_class in TestAgentError.agent_exceptions:
+            exception_instance = exception_class("A test Message")
+
+            self.assertEqual(str(exception_instance), "[{0}] A test Message".format(exception_name))
+
+    def test_agent_exceptions_should_include_the_inner_exception_in_their_error_message(self):
+        inner_exception = Exception("The inner exception")
+
+        for exception_name, exception_class in TestAgentError.agent_exceptions:
+            exception_instance = exception_class("A test Message", inner_exception)
+
+            self.assertEqual(str(exception_instance), "[{0}] A test Message\nInner error: The inner exception".format(exception_name))
diff --git a/tests/common/test_logger.py b/tests/common/test_logger.py
index 2f745f99..116684d3 100644
--- a/tests/common/test_logger.py
+++ b/tests/common/test_logger.py
@@ -18,11 +18,9 @@
 import json
 from datetime import datetime
 
-from mock import patch, MagicMock
-
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.event import add_log_event
-from tests.tools import AgentTestCase
+from tests.tools import AgentTestCase, patch, MagicMock
 
 _MSG_INFO = "This is our test info logging message {0} {1}"
 _MSG_WARN = "This is our test warn logging message {0} {1}"
diff --git a/tests/common/test_version.py b/tests/common/test_version.py
index 14c16a86..6a5c56cd 100644
--- a/tests/common/test_version.py
+++ b/tests/common/test_version.py
@@ -17,14 +17,15 @@
 
 from __future__ import print_function
 
-import textwrap
-
 import mock
+import os
+import textwrap
 
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.version import set_current_agent, \
     AGENT_LONG_VERSION, AGENT_VERSION, AGENT_NAME, AGENT_NAME_PATTERN, \
     get_f5_platform, get_distro
-from tests.tools import *
+from tests.tools import AgentTestCase, open_patch, patch
 
 
 def freebsd_system():
diff --git a/tests/daemon/test_daemon.py b/tests/daemon/test_daemon.py
index 4e87e5a0..91b68bd3 100644
--- a/tests/daemon/test_daemon.py
+++ b/tests/daemon/test_daemon.py
@@ -14,12 +14,17 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 #
+
+import os
+import unittest
+
 from multiprocessing import Process
 
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.daemon import *
 from azurelinuxagent.daemon.main import OPENSSL_FIPS_ENVIRONMENT
 from azurelinuxagent.pa.provision.default import ProvisionHandler
-from tests.tools import *
+from tests.tools import AgentTestCase, Mock, patch
 
 
 class MockDaemonCall(object):
diff --git a/tests/daemon/test_resourcedisk.py b/tests/daemon/test_resourcedisk.py
index 836d7b1c..3f43c733 100644
--- a/tests/daemon/test_resourcedisk.py
+++ b/tests/daemon/test_resourcedisk.py
@@ -15,11 +15,12 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from tests.tools import *
-from azurelinuxagent.common.exception import *
-from azurelinuxagent.daemon import *
+import unittest
+
+from tests.tools import AgentTestCase
 from azurelinuxagent.daemon.resourcedisk.default import ResourceDiskHandler
 
+
 class TestResourceDisk(AgentTestCase):
     def test_mount_flags_empty(self):
         partition = '/dev/sdb1'
diff --git a/tests/distro/test_resourceDisk.py b/tests/distro/test_resourceDisk.py
index 5f9db0a3..32358992 100644
--- a/tests/distro/test_resourceDisk.py
+++ b/tests/distro/test_resourceDisk.py
@@ -21,9 +21,10 @@
 import os
 import stat
 import sys
+import unittest
 from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.daemon.resourcedisk import get_resourcedisk_handler
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestResourceDisk(AgentTestCase):
diff --git a/tests/distro/test_scvmm.py b/tests/distro/test_scvmm.py
index fdca1e5d..a0743596 100644
--- a/tests/distro/test_scvmm.py
+++ b/tests/distro/test_scvmm.py
@@ -19,12 +19,14 @@
 # http://msdn.microsoft.com/en-us/library/cc227259%28PROT.13%29.aspx
 
 import mock
-from tests.tools import *
+from tests.tools import AgentTestCase, Mock, patch
+import unittest
 
 import azurelinuxagent.daemon.scvmm as scvmm
 from azurelinuxagent.daemon.main import *
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
 
+
 class TestSCVMM(AgentTestCase):
     def test_scvmm_detection_with_file(self):
         # setup
diff --git a/tests/ga/test_env.py b/tests/ga/test_env.py
index b69303d9..c0dc7cf1 100644
--- a/tests/ga/test_env.py
+++ b/tests/ga/test_env.py
@@ -16,7 +16,7 @@
 #
 from azurelinuxagent.common.osutil.default import DefaultOSUtil, shellutil
 from azurelinuxagent.ga.env import EnvHandler
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestEnvHandler(AgentTestCase):
@@ -34,7 +34,6 @@ class TestEnvHandler(AgentTestCase):
         self.mock_get_osutil = patch("azurelinuxagent.common.osutil.factory._get_osutil", return_value=DefaultOSUtil())
         self.mock_get_osutil.start()
 
-
     def tearDown(self):
         self.mock_get_osutil.stop()
         AgentTestCase.tearDown(self)
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 2644d59a..0efc6de9 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -16,13 +16,17 @@
 #
 
 import os.path
+import unittest
 
 from datetime import timedelta
 
 from azurelinuxagent.ga.monitor import get_monitor_handler
 from nose.plugins.attrib import attr
-from tests.protocol.mockwiredata import *
+from tests.protocol import mockwiredata
+from tests.tools import are_cgroups_enabled, AgentTestCase, data_dir, i_am_root, MagicMock, Mock, \
+    skip_if_predicate_false, patch
 
+from azurelinuxagent.common.exception import ResourceGoneError
 from azurelinuxagent.common.protocol.restapi import Extension, ExtHandlerProperties
 from azurelinuxagent.ga.exthandlers import *
 from azurelinuxagent.common.protocol.wire import WireProtocol, InVMArtifactsProfile
@@ -360,7 +364,7 @@ class TestExtension(ExtensionTestCase):
         :param args: Any additional args passed to the function, needed for creating a mock for handler and protocol
         :return: test_data, exthandlers_handler, protocol
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install and enable is successful
@@ -381,7 +385,7 @@ class TestExtension(ExtensionTestCase):
         return test_data, exthandlers_handler, protocol
 
     def test_ext_handler(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -444,7 +448,7 @@ class TestExtension(ExtensionTestCase):
         self._assert_no_handler_status(protocol.report_vm_status)
 
     def test_ext_zip_file_packages_removed_in_update_case(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -483,7 +487,7 @@ class TestExtension(ExtensionTestCase):
                                          extension_version="1.2.0")
 
     def test_ext_zip_file_packages_removed_in_uninstall_case(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         extension_version = "1.0.0"
 
@@ -504,7 +508,7 @@ class TestExtension(ExtensionTestCase):
                                          extension_version=extension_version)
 
     def test_ext_zip_file_packages_removed_in_update_and_uninstall_case(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -552,21 +556,21 @@ class TestExtension(ExtensionTestCase):
                                          extension_version="1.2.0")
 
     def test_ext_handler_no_settings(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_NO_SETTINGS)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_NO_SETTINGS)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         exthandlers_handler.run()
         self._assert_handler_status(protocol.report_vm_status, "Ready", 0, "1.0.0")
 
     def test_ext_handler_no_public_settings(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_NO_PUBLIC)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_NO_PUBLIC)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         exthandlers_handler.run()
         self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.0.0")
 
     def test_ext_handler_no_ext(self, *args):
-        test_data = WireProtocolData(DATA_FILE_NO_EXT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_NO_EXT)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Assert no extension handler status
@@ -574,7 +578,7 @@ class TestExtension(ExtensionTestCase):
         self._assert_no_handler_status(protocol.report_vm_status)
 
     def test_ext_handler_sequencing(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_SEQUENCING)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SEQUENCING)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -647,14 +651,14 @@ class TestExtension(ExtensionTestCase):
         self.assertEqual(exthandlers_handler.ext_handlers.extHandlers[1].properties.extensions[0].dependencyLevel, 5)
 
     def test_ext_handler_sequencing_default_dependency_level(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.run()
         self.assertEqual(exthandlers_handler.ext_handlers.extHandlers[0].properties.extensions[0].dependencyLevel, 0)
         self.assertEqual(exthandlers_handler.ext_handlers.extHandlers[0].properties.extensions[0].dependencyLevel, 0)
 
     def test_ext_handler_sequencing_invalid_dependency_level(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_SEQUENCING)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SEQUENCING)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         test_data.goal_state = test_data.goal_state.replace("<Incarnation>1<",
@@ -716,7 +720,7 @@ class TestExtension(ExtensionTestCase):
 
         expected_status_json = json.loads(expected_status)
 
-        test_data = WireProtocolData(DATA_FILE_MULTIPLE_EXT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_MULTIPLE_EXT)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.run()
 
@@ -726,7 +730,7 @@ class TestExtension(ExtensionTestCase):
         self.assertEquals(expected_status_json, actual_status_json)
 
     def test_ext_handler_rollingupgrade(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_ROLLINGUPGRADE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_ROLLINGUPGRADE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -806,7 +810,7 @@ class TestExtension(ExtensionTestCase):
     def test_ext_handler_download_failure_transient(self, mock_add_event, *args):
         original_sleep = time.sleep
 
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.download_ext_handler_pkg = Mock(side_effect=ProtocolError)
 
@@ -817,7 +821,7 @@ class TestExtension(ExtensionTestCase):
     @patch('azurelinuxagent.common.errorstate.ErrorState.is_triggered')
     @patch('azurelinuxagent.ga.exthandlers.add_event')
     def test_ext_handler_report_status_permanent(self, mock_add_event, mock_error_state, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.report_vm_status = Mock(side_effect=ProtocolError)
 
@@ -831,7 +835,7 @@ class TestExtension(ExtensionTestCase):
 
     @patch('azurelinuxagent.ga.exthandlers.add_event')
     def test_ext_handler_report_status_resource_gone(self, mock_add_event, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.report_vm_status = Mock(side_effect=ResourceGoneError)
 
@@ -845,7 +849,7 @@ class TestExtension(ExtensionTestCase):
     @patch('azurelinuxagent.common.errorstate.ErrorState.is_triggered')
     @patch('azurelinuxagent.ga.exthandlers.ExtHandlerInstance.report_event')
     def test_ext_handler_download_failure_permanent_ProtocolError(self, mock_add_event, mock_error_state, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.get_ext_handler_pkgs = Mock(side_effect=ProtocolError)
 
@@ -863,7 +867,7 @@ class TestExtension(ExtensionTestCase):
     @patch('azurelinuxagent.common.event.add_event')
     def test_ext_handler_download_failure_permanent_with_ExtensionDownloadError_and_triggered(self, mock_add_event,
                                                                                               mock_error_state, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.get_ext_handler_pkgs = Mock(side_effect=ExtensionDownloadError)
 
@@ -883,7 +887,7 @@ class TestExtension(ExtensionTestCase):
     def test_ext_handler_download_failure_permanent_with_ExtensionDownloadError_and_not_triggered(self, mock_add_event,
                                                                                                   mock_error_state,
                                                                                                   *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         protocol.get_ext_handler_pkgs = Mock(side_effect=ExtensionDownloadError)
 
@@ -895,7 +899,7 @@ class TestExtension(ExtensionTestCase):
 
     @patch('azurelinuxagent.ga.exthandlers.fileutil')
     def test_ext_handler_io_error(self, mock_fileutil, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         mock_fileutil.write_file.return_value = IOError("Mock IO Error")
@@ -932,7 +936,7 @@ class TestExtension(ExtensionTestCase):
                         self.assertTrue(exthandlers_handler.extension_processing_allowed())
 
     def test_handle_ext_handlers_on_hold_true(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.ext_handlers, exthandlers_handler.last_etag = protocol.get_ext_handlers()
         protocol.get_artifacts_profile = MagicMock()
@@ -951,7 +955,7 @@ class TestExtension(ExtensionTestCase):
             self.assertEqual(1, patch_handle_ext_handlers.call_count)
 
     def test_handle_ext_handlers_on_hold_false(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.ext_handlers, exthandlers_handler.last_etag = protocol.get_ext_handlers()
         exthandlers_handler.protocol = protocol
@@ -975,7 +979,7 @@ class TestExtension(ExtensionTestCase):
             self.assertEqual(1, patch_handle_ext_handler.call_count)
 
     def test_last_etag_on_extension_processing(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.ext_handlers, etag = protocol.get_ext_handlers()
         exthandlers_handler.protocol = protocol
@@ -1002,7 +1006,7 @@ class TestExtension(ExtensionTestCase):
         self.assertEquals(expected_seq_no, ext_status.sequenceNumber)
 
     def test_ext_handler_no_reporting_status(self, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.run()
         self._assert_handler_status(protocol.report_vm_status, "Ready", 1, "1.0.0")
@@ -1023,7 +1027,7 @@ class TestExtension(ExtensionTestCase):
         Testing wait_for_handler_successful_completion() when there is no extension in a handler.
         Expected to return True.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         handler = ExtHandler(name="handler")
@@ -1051,7 +1055,7 @@ class TestExtension(ExtensionTestCase):
         Testing wait_for_handler_successful_completion() when there is no status file or seq_no is negative.
         Expected to return False.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         ExtHandlerInstance.get_ext_handling_status = MagicMock(return_value=None)
@@ -1062,7 +1066,7 @@ class TestExtension(ExtensionTestCase):
         Testing wait_for_handler_successful_completion() when there is successful status.
         Expected to return True.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         status = "success"
@@ -1075,7 +1079,7 @@ class TestExtension(ExtensionTestCase):
         Testing wait_for_handler_successful_completion() when there is error status.
         Expected to return False.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         status = "error"
@@ -1088,7 +1092,7 @@ class TestExtension(ExtensionTestCase):
         Testing wait_for_handler_successful_completion() when there is non terminal status.
         Expected to return False.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Choose a non-terminal status
@@ -1102,7 +1106,7 @@ class TestExtension(ExtensionTestCase):
         Testing get_ext_handling_status() function with various cases and
         verifying against the expected values
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         handler_name = "Handler"
@@ -1145,7 +1149,7 @@ class TestExtension(ExtensionTestCase):
         Testing is_ext_handling_complete() with various input and
         verifying against the expected output values.
         '''
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         handler_name = "Handler"
@@ -1183,18 +1187,18 @@ class TestExtension(ExtensionTestCase):
                     config_version = '1.3.0'
                     decision_version = '1.3.0'
                     if autoupgrade:
-                        datafile = DATA_FILE_EXT_AUTOUPGRADE_INTERNALVERSION
+                        datafile = mockwiredata.DATA_FILE_EXT_AUTOUPGRADE_INTERNALVERSION
                     else:
-                        datafile = DATA_FILE_EXT_INTERNALVERSION
+                        datafile = mockwiredata.DATA_FILE_EXT_INTERNALVERSION
                 else:
                     config_version = '1.0.0'
                     decision_version = '1.0.0'
                     if autoupgrade:
-                        datafile = DATA_FILE_EXT_AUTOUPGRADE
+                        datafile = mockwiredata.DATA_FILE_EXT_AUTOUPGRADE
                     else:
-                        datafile = DATA_FILE
+                        datafile = mockwiredata.DATA_FILE
 
-                _, protocol = self._create_mock(WireProtocolData(datafile), *args)
+                _, protocol = self._create_mock(mockwiredata.WireProtocolData(datafile), *args)
                 ext_handlers, _ = protocol.get_ext_handlers()
                 self.assertEqual(1, len(ext_handlers.extHandlers))
                 ext_handler = ext_handlers.extHandlers[0]
@@ -1227,7 +1231,7 @@ class TestExtension(ExtensionTestCase):
             (None, '4.1', '4.1.0.0'),
         ]
 
-        _, protocol = self._create_mock(WireProtocolData(DATA_FILE), *args)
+        _, protocol = self._create_mock(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE), *args)
         version_uri = Mock()
         version_uri.uri = 'http://some/Microsoft.OSTCExtensions_ExampleHandlerLinux_asiaeast_manifest.xml'
 
@@ -1247,19 +1251,19 @@ class TestExtension(ExtensionTestCase):
     @patch('azurelinuxagent.common.conf.get_extensions_enabled', return_value=False)
     def test_extensions_disabled(self, _, *args):
         # test status is reported for no extensions
-        test_data = WireProtocolData(DATA_FILE_NO_EXT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_NO_EXT)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.run()
         self._assert_no_handler_status(protocol.report_vm_status)
 
         # test status is reported, but extensions are not processed
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         exthandlers_handler.run()
         self._assert_no_handler_status(protocol.report_vm_status)
 
     def test_extensions_deleted(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_DELETION)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_DELETION)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial enable is successful
@@ -1286,7 +1290,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension install fails, the operation should not be retried.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install is unsuccessful
@@ -1309,7 +1313,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension install fails, the operation should be reported to our telemetry service.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure install is unsuccessful
@@ -1324,7 +1328,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension enable fails, the operation should not be retried.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install is successful, but enable fails
@@ -1347,7 +1351,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension enable fails, the operation should be reported.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install is successful, but enable fails
@@ -1364,7 +1368,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension disable fails, the operation should not be retried.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install and enable is successful, but disable fails
@@ -1400,7 +1404,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension disable fails, the operation should be reported.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install and enable is successful, but disable fails
@@ -1428,7 +1432,7 @@ class TestExtension(ExtensionTestCase):
         """
         When extension uninstall fails, the operation should not be retried.
         """
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         # Ensure initial install and enable is successful, but uninstall fails
@@ -1743,7 +1747,7 @@ class TestExtension(ExtensionTestCase):
         self._assert_ext_status(protocol.report_ext_status, "success", 0)
 
     def test_ext_path_and_version_env_variables_set_for_ever_operation(self, *args):
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         with patch.object(CGroupConfigurator.get_instance(), "start_extension_command") as patch_start_cmd:
@@ -1761,7 +1765,7 @@ class TestExtension(ExtensionTestCase):
 
     @patch("azurelinuxagent.common.cgroupconfigurator.handle_process_completion", side_effect="Process Successful")
     def test_ext_sequence_no_should_be_set_for_every_command_call(self, _, *args):
-        test_data = WireProtocolData(DATA_FILE_MULTIPLE_EXT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_MULTIPLE_EXT)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
 
         with patch("subprocess.Popen") as patch_popen:
@@ -1812,7 +1816,7 @@ class TestExtension(ExtensionTestCase):
         base_dir = os.path.join(conf.get_lib_dir(), 'OSTCExtensions.ExampleHandlerLinux-1.0.0', test_file_name)
         self.create_script(test_file_name, test_file, base_dir)
 
-        test_data = WireProtocolData(DATA_FILE_EXT_SINGLE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_SINGLE)
         exthandlers_handler, protocol = self._create_mock(test_data, *args)
         expected_seq_no = 0
 
@@ -1899,7 +1903,7 @@ class TestExtension(ExtensionTestCase):
 class TestExtensionSequencing(AgentTestCase):
 
     def _create_mock(self, mock_http_get, MockCryptUtil):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
 
         # Mock protocol to return test data
         mock_http_get.side_effect = test_data.mock_http_get
@@ -2156,7 +2160,7 @@ class TestExtensionWithCGroupsEnabled(AgentTestCase):
     def test_ext_handler_with_cgroup_enabled(self, *args):
         self.assertTrue(i_am_root(), "Test does not run when non-root")
 
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, _, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
@@ -2223,7 +2227,7 @@ class TestExtensionWithCGroupsEnabled(AgentTestCase):
     def test_ext_handler_and_monitor_handler_with_cgroup_enabled(self, patch_add_event, *args):
         self.assertTrue(i_am_root(), "Test does not run when non-root")
 
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, monitor_handler, protocol = self._create_mock(test_data, *args)
 
         monitor_handler.last_cgroup_polling_telemetry = datetime.datetime.utcnow() - timedelta(hours=1)
@@ -2258,7 +2262,7 @@ class TestExtensionWithCGroupsEnabled(AgentTestCase):
         from azurelinuxagent.common.cgroupapi import CGroupsApi
         print(CGroupsApi._is_systemd())
 
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         exthandlers_handler, _, protocol = self._create_mock(test_data, *args)
 
         # Test enable scenario.
diff --git a/tests/ga/test_exthandlers.py b/tests/ga/test_exthandlers.py
index d54843d2..76e84ea1 100644
--- a/tests/ga/test_exthandlers.py
+++ b/tests/ga/test_exthandlers.py
@@ -2,14 +2,17 @@
 # Licensed under the Apache License.
 import json
 import subprocess
+import os
+import time
 
 from azurelinuxagent.common.protocol.restapi import ExtensionStatus, Extension, ExtHandler, ExtHandlerProperties
 from azurelinuxagent.ga.exthandlers import parse_ext_status, ExtHandlerInstance, get_exthandlers_handler, \
     ExtCommandEnvVariable
+from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.exception import ProtocolError, ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.event import WALAEventOperation
 from azurelinuxagent.common.utils.extensionprocessutil import TELEMETRY_MESSAGE_MAX_LEN, format_stdout_stderr, read_output
-from tests.tools import *
+from tests.tools import AgentTestCase, patch, mock_sleep
 
 
 class TestExtHandlers(AgentTestCase):
diff --git a/tests/ga/test_exthandlers_download_extension.py b/tests/ga/test_exthandlers_download_extension.py
index 329399cf..dc1e478b 100644
--- a/tests/ga/test_exthandlers_download_extension.py
+++ b/tests/ga/test_exthandlers_download_extension.py
@@ -1,13 +1,13 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the Apache License.
 
-import zipfile, time
+import zipfile, time, os
 
 from azurelinuxagent.common.protocol.restapi import ExtHandler, ExtHandlerProperties, ExtHandlerPackage, ExtHandlerVersionUri
 from azurelinuxagent.common.protocol.wire import WireProtocol
 from azurelinuxagent.ga.exthandlers import ExtHandlerInstance, NUMBER_OF_DOWNLOAD_RETRIES
 from azurelinuxagent.common.exception import ExtensionDownloadError, ExtensionErrorCodes
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class DownloadExtensionTestCase(AgentTestCase):
diff --git a/tests/ga/test_exthandlers_exthandlerinstance.py b/tests/ga/test_exthandlers_exthandlerinstance.py
index 1f392c16..159fefb2 100644
--- a/tests/ga/test_exthandlers_exthandlerinstance.py
+++ b/tests/ga/test_exthandlers_exthandlerinstance.py
@@ -4,7 +4,10 @@
 from azurelinuxagent.ga.exthandlers import ExtHandlerInstance
 from azurelinuxagent.common.protocol.restapi import ExtHandler, ExtHandlerProperties, ExtHandlerPackage, \
     ExtHandlerVersionUri
-from tests.tools import *
+import os
+import shutil
+import sys
+from tests.tools import AgentTestCase, patch
 
 
 class ExtHandlerInstanceTestCase(AgentTestCase):
diff --git a/tests/ga/test_monitor.py b/tests/ga/test_monitor.py
index 217effee..4523694f 100644
--- a/tests/ga/test_monitor.py
+++ b/tests/ga/test_monitor.py
@@ -26,12 +26,12 @@ import tempfile
 import time
 from datetime import timedelta
 
-from mock import Mock, MagicMock, patch
 from nose.plugins.attrib import attr
 
 from azurelinuxagent.common import logger
 from azurelinuxagent.common.cgroup import CGroup
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.datacontract import get_properties
 from azurelinuxagent.common.event import EventLogger, WALAEventOperation, CONTAINER_ID_ENV_VARIABLE
 from azurelinuxagent.common.exception import HttpError
@@ -45,8 +45,9 @@ from azurelinuxagent.common.version import AGENT_VERSION, CURRENT_VERSION, AGENT
 from azurelinuxagent.ga.monitor import parse_xml_event, get_monitor_handler, MonitorHandler, \
     generate_extension_metrics_telemetry_dictionary, parse_json_event
 from tests.common.test_cgroupstelemetry import make_new_cgroup, consume_cpu_time, consume_memory
-from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE, conf
-from tests.tools import load_data, AgentTestCase, data_dir, are_cgroups_enabled, i_am_root, skip_if_predicate_false
+from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE
+from tests.tools import Mock, MagicMock, patch, load_data, AgentTestCase, data_dir, are_cgroups_enabled, \
+    i_am_root, skip_if_predicate_false
 
 
 class ResponseMock(Mock):
diff --git a/tests/ga/test_remoteaccess.py b/tests/ga/test_remoteaccess.py
index 87eb64c5..c0d0cfc6 100644
--- a/tests/ga/test_remoteaccess.py
+++ b/tests/ga/test_remoteaccess.py
@@ -16,9 +16,9 @@
 #
 import xml
 
-from tests.tools import *
+from tests.tools import AgentTestCase, load_data, patch
 from azurelinuxagent.common.protocol.wire import *
-from azurelinuxagent.common.osutil import get_osutil
+
 
 class TestRemoteAccess(AgentTestCase):
     def test_parse_remote_access(self):
diff --git a/tests/ga/test_remoteaccess_handler.py b/tests/ga/test_remoteaccess_handler.py
index 71271c97..a54e4d54 100644
--- a/tests/ga/test_remoteaccess_handler.py
+++ b/tests/ga/test_remoteaccess_handler.py
@@ -21,7 +21,7 @@ from azurelinuxagent.common.exception import RemoteAccessError
 from azurelinuxagent.common.protocol.wire import *
 from azurelinuxagent.ga.remoteaccess import RemoteAccessHandler
 from tests.common.osutil.mock_osutil import MockOSUtil
-from tests.tools import *
+from tests.tools import AgentTestCase, load_data, patch
 
 
 info_messages = []
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index c435dcc7..38cb8712 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -3,14 +3,15 @@
 
 from __future__ import print_function
 
-from azurelinuxagent.common.event import *
+import tempfile
+import unittest
+
 from azurelinuxagent.common.protocol.hostplugin import *
 from azurelinuxagent.common.protocol.metadata import *
 from azurelinuxagent.common.protocol.wire import *
-from azurelinuxagent.common.utils.fileutil import *
 from azurelinuxagent.ga.update import *
+from tests.tools import AgentTestCase, call, data_dir, DEFAULT, patch, load_bin_data, load_data, Mock, MagicMock
 
-from tests.tools import *
 
 NO_ERROR = {
     "last_failure" : 0.0,
diff --git a/tests/pa/test_deprovision.py b/tests/pa/test_deprovision.py
index f4be81e2..591a2c67 100644
--- a/tests/pa/test_deprovision.py
+++ b/tests/pa/test_deprovision.py
@@ -15,14 +15,15 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-import signal
+import os
 import tempfile
+import unittest
 
 import azurelinuxagent.common.utils.fileutil as fileutil
 
 from azurelinuxagent.pa.deprovision import get_deprovision_handler
 from azurelinuxagent.pa.deprovision.default import DeprovisionHandler
-from tests.tools import *
+from tests.tools import AgentTestCase, distros, Mock, patch
 
 
 class TestDeprovision(AgentTestCase):
diff --git a/tests/pa/test_provision.py b/tests/pa/test_provision.py
index 54c2f8c8..54175601 100644
--- a/tests/pa/test_provision.py
+++ b/tests/pa/test_provision.py
@@ -15,13 +15,20 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import os
+import re
+import tempfile
+import unittest
+
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import ProvisionError
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
 from azurelinuxagent.common.protocol import OVF_FILE_NAME
 from azurelinuxagent.pa.provision import get_provision_handler
 from azurelinuxagent.pa.provision.cloudinit import CloudInitProvisionHandler
 from azurelinuxagent.pa.provision.default import ProvisionHandler
-from tests.tools import *
+from azurelinuxagent.common.utils import fileutil
+from tests.tools import AgentTestCase, distros, load_data, MagicMock, Mock, patch
 
 
 class TestProvision(AgentTestCase):
@@ -372,6 +379,7 @@ class TestProvision(AgentTestCase):
         provisioning_handler = get_provision_handler()
         self.assertIsInstance(provisioning_handler, CloudInitProvisionHandler, 'Provisioning handler should be cloud-init if agent is set to cloud-init')
 
+
 if __name__ == '__main__':
     unittest.main()
 
diff --git a/tests/protocol/mockmetadata.py b/tests/protocol/mockmetadata.py
index 65399a85..6f6de003 100644
--- a/tests/protocol/mockmetadata.py
+++ b/tests/protocol/mockmetadata.py
@@ -15,9 +15,8 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from tests.tools import *
+from tests.tools import load_data, MagicMock, Mock
 from azurelinuxagent.common.future import httpclient
-from azurelinuxagent.common.utils.cryptutil import CryptUtil
 
 DATA_FILE = {
         "identity": "metadata/identity.json",
@@ -32,6 +31,7 @@ DATA_FILE = {
 DATA_FILE_NO_EXT = DATA_FILE.copy()
 DATA_FILE_NO_EXT["ext_handlers"] = "metadata/ext_handlers_no_ext.json"
 
+
 class MetadataProtocolData(object):
     def __init__(self, data_files):
         self.identity = load_data(data_files.get("identity"))
diff --git a/tests/protocol/mockwiredata.py b/tests/protocol/mockwiredata.py
index 2adf8c89..5202003e 100644
--- a/tests/protocol/mockwiredata.py
+++ b/tests/protocol/mockwiredata.py
@@ -15,7 +15,7 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from tests.tools import *
+from tests.tools import load_bin_data, load_data, MagicMock, Mock
 from azurelinuxagent.common.exception import HttpError, ResourceGoneError
 from azurelinuxagent.common.future import httpclient
 from azurelinuxagent.common.utils.cryptutil import CryptUtil
diff --git a/tests/protocol/test_healthservice.py b/tests/protocol/test_healthservice.py
index 1bd44adb..8af322f3 100644
--- a/tests/protocol/test_healthservice.py
+++ b/tests/protocol/test_healthservice.py
@@ -19,7 +19,7 @@ from azurelinuxagent.common.exception import HttpError
 from azurelinuxagent.common.protocol.healthservice import Observation, HealthService
 from azurelinuxagent.common.utils import restutil
 from tests.protocol.test_hostplugin import MockResponse
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestHealthService(AgentTestCase):
diff --git a/tests/protocol/test_hostplugin.py b/tests/protocol/test_hostplugin.py
index 842ef03e..3aa92669 100644
--- a/tests/protocol/test_hostplugin.py
+++ b/tests/protocol/test_hostplugin.py
@@ -19,6 +19,7 @@ import base64
 import json
 import sys
 import datetime
+import unittest
 
 import azurelinuxagent.common.protocol.restapi as restapi
 import azurelinuxagent.common.protocol.wire as wire
@@ -31,7 +32,7 @@ from azurelinuxagent.common.protocol.hostplugin import API_VERSION
 from azurelinuxagent.common.utils import restutil
 from tests.protocol.mockwiredata import WireProtocolData, DATA_FILE
 from tests.protocol.test_wire import MockResponse
-from tests.tools import *
+from tests.tools import AgentTestCase, PY_VERSION_MAJOR, Mock, patch
 
 if sys.version_info[0] == 3:
     import http.client as httpclient
diff --git a/tests/protocol/test_image_info_matcher.py b/tests/protocol/test_image_info_matcher.py
index 936e94ff..bd538b39 100644
--- a/tests/protocol/test_image_info_matcher.py
+++ b/tests/protocol/test_image_info_matcher.py
@@ -16,9 +16,8 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 
-from azurelinuxagent.common.datacontract import set_properties
 from azurelinuxagent.common.protocol.imds import ImageInfoMatcher
-from tests.tools import *
+import unittest
 
 
 class TestImageInfoMatcher(unittest.TestCase):
diff --git a/tests/protocol/test_imds.py b/tests/protocol/test_imds.py
index 2ac172d6..68c3d436 100644
--- a/tests/protocol/test_imds.py
+++ b/tests/protocol/test_imds.py
@@ -17,6 +17,8 @@
 # Requires Python 2.6+ and Openssl 1.0+
 
 import json
+import os
+import unittest
 
 import azurelinuxagent.common.protocol.imds as imds
 
@@ -25,7 +27,7 @@ from azurelinuxagent.common.exception import HttpError, ResourceGoneError
 from azurelinuxagent.common.future import ustr, httpclient
 from azurelinuxagent.common.utils import restutil
 from tests.ga.test_update import ResponseMock
-from tests.tools import *
+from tests.tools import AgentTestCase, data_dir, MagicMock, Mock, patch
 
 
 class TestImds(AgentTestCase):
diff --git a/tests/protocol/test_metadata.py b/tests/protocol/test_metadata.py
index 3f1793ae..d5ec36de 100644
--- a/tests/protocol/test_metadata.py
+++ b/tests/protocol/test_metadata.py
@@ -15,15 +15,13 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from azurelinuxagent.common.datacontract import get_properties, set_properties
-from azurelinuxagent.common.exception import ProtocolError
 from azurelinuxagent.common.protocol.metadata import *
 from azurelinuxagent.common.protocol.restapi import *
 from azurelinuxagent.common.telemetryevent import TelemetryEventList, TelemetryEvent
 from azurelinuxagent.common.utils import restutil
 
 from tests.protocol.mockmetadata import *
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestMetadataProtocolGetters(AgentTestCase):
diff --git a/tests/protocol/test_protocol_util.py b/tests/protocol/test_protocol_util.py
index 4a117899..6f341f10 100644
--- a/tests/protocol/test_protocol_util.py
+++ b/tests/protocol/test_protocol_util.py
@@ -15,10 +15,17 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from tests.tools import *
+import os
+import unittest
+
+from tests.tools import AgentTestCase, MagicMock, Mock, patch
 from azurelinuxagent.common.exception import *
 from azurelinuxagent.common.protocol import get_protocol_util, \
                                             TAG_FILE_NAME
+from azurelinuxagent.common.utils.restutil import KNOWN_WIRESERVER_IP
+from azurelinuxagent.common.protocol.util import ENDPOINT_FILE_NAME
+from errno import ENOENT
+
 
 @patch("time.sleep")
 class TestProtocolUtil(AgentTestCase):
@@ -34,22 +41,22 @@ class TestProtocolUtil(AgentTestCase):
         protocol_util.dhcp_handler = MagicMock()
         protocol_util.dhcp_handler.endpoint = "foo.bar"
 
-        #Test wire protocol is available
+        # Test wire protocol is available
         protocol = protocol_util.get_protocol()
         self.assertEquals(WireProtocol.return_value, protocol)
 
-        #Test wire protocol is not available
+        # Test wire protocol is not available
         protocol_util.clear_protocol()
         WireProtocol.return_value.detect.side_effect = ProtocolError()
 
         protocol = protocol_util.get_protocol()
         self.assertEquals(MetadataProtocol.return_value, protocol)
 
-        #Test no protocol is available
+        # Test no protocol is available
         protocol_util.clear_protocol()
         WireProtocol.return_value.detect.side_effect = ProtocolError()
-
         MetadataProtocol.return_value.detect.side_effect = ProtocolError()
+
         self.assertRaises(ProtocolError, protocol_util.get_protocol)
 
     def test_detect_protocol_by_file(self, _):
@@ -59,12 +66,12 @@ class TestProtocolUtil(AgentTestCase):
 
         tag_file = os.path.join(self.tmp_dir, TAG_FILE_NAME)
 
-        #Test tag file doesn't exist
+        # Test tag file doesn't exist
         protocol_util.get_protocol(by_file=True)
         protocol_util._detect_wire_protocol.assert_any_call()
         protocol_util._detect_metadata_protocol.assert_not_called()
 
-        #Test tag file exists
+        # Test tag file exists
         protocol_util.clear_protocol()
         protocol_util._detect_wire_protocol.reset_mock()
         protocol_util._detect_metadata_protocol.reset_mock()
@@ -75,6 +82,156 @@ class TestProtocolUtil(AgentTestCase):
         protocol_util._detect_metadata_protocol.assert_any_call()
         protocol_util._detect_wire_protocol.assert_not_called()
 
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    @patch("azurelinuxagent.common.protocol.util.WireProtocol")
+    def test_detect_wire_protocol_no_dhcp(self, WireProtocol, mock_get_lib_dir, _):
+        WireProtocol.return_value.detect = Mock()
+        mock_get_lib_dir.return_value = self.tmp_dir
+
+        protocol_util = get_protocol_util()
+
+        protocol_util.osutil = MagicMock()
+        protocol_util.osutil.is_dhcp_available.return_value = False
+
+        protocol_util.dhcp_handler = MagicMock()
+        protocol_util.dhcp_handler.endpoint = None
+        protocol_util.dhcp_handler.run = Mock()
+
+        endpoint_file = protocol_util._get_wireserver_endpoint_file_path()
+
+        # Test wire protocol when no endpoint file has been written
+        protocol_util._detect_wire_protocol()
+        self.assertEqual(KNOWN_WIRESERVER_IP, protocol_util.get_wireserver_endpoint())
+
+        # Test wire protocol when endpoint was previously detected
+        protocol_util.clear_protocol()
+        with open(endpoint_file, "w+") as endpoint_fd:
+            endpoint_fd.write("baz.qux")
+
+        protocol_util._detect_wire_protocol()
+        self.assertEqual("baz.qux", protocol_util.get_wireserver_endpoint())
+
+        # Test wire protocol on dhcp failure
+        protocol_util.clear_protocol()
+        protocol_util.osutil.is_dhcp_available.return_value = True
+        protocol_util.dhcp_handler.run.side_effect = DhcpError()
+
+        self.assertRaises(ProtocolError, protocol_util._detect_wire_protocol)
+
+    @patch("azurelinuxagent.common.protocol.util.MetadataProtocol")
+    @patch("azurelinuxagent.common.protocol.util.WireProtocol")
+    def test_get_protocol(self, WireProtocol, MetadataProtocol, _):
+        WireProtocol.return_value = MagicMock()
+        MetadataProtocol.return_value = MagicMock()
+
+        protocol_util = get_protocol_util()
+        protocol_util.get_wireserver_endpoint = Mock()
+        protocol_util._detect_protocol = MagicMock()
+
+        # Test for wire protocol
+        protocol_util._save_protocol("WireProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(WireProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_any_call()
+
+        # Test to ensure protocol persists
+        protocol_util.get_wireserver_endpoint.reset_mock()
+        protocol_util._save_protocol("MetadataProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(WireProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+        # Test for metadata protocol
+        protocol_util.clear_protocol()
+        protocol_util._save_protocol("MetadataProtocol")
+
+        protocol = protocol_util.get_protocol()
+        self.assertEquals(MetadataProtocol.return_value, protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+        # Test for unknown protocol
+        protocol_util.clear_protocol()
+        protocol_util._save_protocol("Not_a_Protocol")
+        protocol_util._detect_protocol.side_effect = NotImplementedError()
+
+        self.assertRaises(NotImplementedError, protocol_util.get_protocol)
+        protocol_util.get_wireserver_endpoint.assert_not_called()
+
+    @patch("azurelinuxagent.common.utils.fileutil")
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    def test_endpoint_file_states(self, mock_get_lib_dir, mock_fileutil, _):
+        mock_get_lib_dir.return_value = self.tmp_dir
+        mock_fileutil = MagicMock()
+
+        protocol_util = get_protocol_util()
+        endpoint_file = protocol_util._get_wireserver_endpoint_file_path()
+
+        # Test get endpoint for io error
+        mock_fileutil.read_file.side_effect = IOError()
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test get endpoint when file not found
+        mock_fileutil.read_file.side_effect = IOError(ENOENT, 'File not found')
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test get endpoint for empty file
+        mock_fileutil.read_file.return_value = ""
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertEquals(ep, KNOWN_WIRESERVER_IP)
+
+        # Test set endpoint for io error
+        mock_fileutil.write_file.side_effect = IOError()
+
+        ep = protocol_util.get_wireserver_endpoint()
+        self.assertRaises(OSUtilError, protocol_util._set_wireserver_endpoint('abc'))
+
+        # Test clear endpoint for io error
+        with open(endpoint_file, "w+") as ep_fd:
+            ep_fd.write("")
+
+        with patch('os.remove') as mock_remove:
+            protocol_util._clear_wireserver_endpoint()
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(endpoint_file, mock_remove.call_args_list[0][0][0])
+
+        # Test clear endpoint when file not found
+        with patch('os.remove') as mock_remove:
+            mock_remove = Mock(side_effect=IOError(ENOENT, 'File not found'))
+            protocol_util._clear_wireserver_endpoint()
+            mock_remove.assert_not_called()
+
+    def test_protocol_file_states(self, _):
+        protocol_util = get_protocol_util()
+        protocol_util._clear_wireserver_endpoint = Mock()
+
+        protocol_file = protocol_util._get_protocol_file_path()
+
+        # Test clear protocol for io error
+        with open(protocol_file, "w+") as proto_fd:
+            proto_fd.write("")
+
+        with patch('os.remove') as mock_remove:
+            protocol_util.clear_protocol()
+            self.assertEqual(1, protocol_util._clear_wireserver_endpoint.call_count)
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(protocol_file, mock_remove.call_args_list[0][0][0])
+
+        # Test clear protocol when file not found
+        protocol_util._clear_wireserver_endpoint.reset_mock()
+
+        with patch('os.remove') as mock_remove:
+            protocol_util.clear_protocol()
+            self.assertEqual(1, protocol_util._clear_wireserver_endpoint.call_count)
+            self.assertEqual(1, mock_remove.call_count)
+            self.assertEqual(protocol_file, mock_remove.call_args_list[0][0][0])
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index b7a324a7..71cc9692 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -16,13 +16,17 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import stat
+import tempfile
+import unittest
 import zipfile
 
 from azurelinuxagent.common.telemetryevent import TelemetryEvent, TelemetryEventParam
 from azurelinuxagent.common.protocol.wire import *
 from azurelinuxagent.common.utils.shellutil import run_get_output
 from tests.ga.test_monitor import random_generator
-from tests.protocol.mockwiredata import *
+from tests.protocol import mockwiredata
+from tests.tools import ANY, AgentTestCase, MagicMock, Mock, patch, running_under_travis, skip_if_predicate_true
 
 data_with_bom = b'\xef\xbb\xbfhehe'
 testurl = 'http://foo'
@@ -83,37 +87,37 @@ class TestWireProtocol(AgentTestCase):
 
     def test_getters(self, *args):
         """Normal case"""
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         self._test_getters(test_data, True, *args)
 
     def test_getters_no_ext(self, *args):
         """Provision with agent is not checked"""
-        test_data = WireProtocolData(DATA_FILE_NO_EXT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_NO_EXT)
         self._test_getters(test_data, True, *args)
 
     def test_getters_ext_no_settings(self, *args):
         """Extensions without any settings"""
-        test_data = WireProtocolData(DATA_FILE_EXT_NO_SETTINGS)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_NO_SETTINGS)
         self._test_getters(test_data, True, *args)
 
     def test_getters_ext_no_public(self, *args):
         """Extensions without any public settings"""
-        test_data = WireProtocolData(DATA_FILE_EXT_NO_PUBLIC)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_EXT_NO_PUBLIC)
         self._test_getters(test_data, True, *args)
 
     def test_getters_ext_no_cert_format(self, *args):
         """Certificate format not specified"""
-        test_data = WireProtocolData(DATA_FILE_NO_CERT_FORMAT)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_NO_CERT_FORMAT)
         self._test_getters(test_data, True, *args)
 
     def test_getters_ext_cert_format_not_pfx(self, *args):
         """Certificate format is not Pkcs7BlobWithPfxContents specified"""
-        test_data = WireProtocolData(DATA_FILE_CERT_FORMAT_NOT_PFX)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE_CERT_FORMAT_NOT_PFX)
         self._test_getters(test_data, False, *args)
 
     @patch("azurelinuxagent.common.protocol.healthservice.HealthService.report_host_plugin_extension_artifact")
     def test_getters_with_stale_goal_state(self, patch_report, *args):
-        test_data = WireProtocolData(DATA_FILE)
+        test_data = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE)
         test_data.emulate_stale_goal_state = True
 
         self._test_getters(test_data, True, *args)
@@ -167,7 +171,7 @@ class TestWireProtocol(AgentTestCase):
 
     def test_status_blob_parsing(self, *args):
         wire_protocol_client = WireProtocol(wireserver_url).client
-        wire_protocol_client.ext_conf = ExtensionsConfig(WireProtocolData(DATA_FILE).ext_conf)
+        wire_protocol_client.ext_conf = ExtensionsConfig(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).ext_conf)
         self.assertEqual(wire_protocol_client.ext_conf.status_upload_blob,
                          u'https://yuezhatest.blob.core.windows.net/vhds/test'
                          u'-cs12.test-cs12.test-cs12.status?sr=b&sp=rw&se'
@@ -179,7 +183,7 @@ class TestWireProtocol(AgentTestCase):
 
     def test_get_host_ga_plugin(self, *args):
         wire_protocol_client = WireProtocol(wireserver_url).client
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
 
         with patch.object(WireClient, "get_goal_state", return_value=goal_state) as patch_get_goal_state:
             host_plugin = wire_protocol_client.get_host_plugin()
@@ -265,7 +269,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client.ext_conf.status_upload_blob = testurl
         wire_protocol_client.ext_conf.status_upload_blob_type = testtype
         wire_protocol_client.status_blob.vm_status = vmstatus
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
 
         with patch.object(HostPluginProtocol,
                           "ensure_initialized",
@@ -311,7 +315,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client.ext_conf.status_upload_blob = testurl
         wire_protocol_client.ext_conf.status_upload_blob_type = testtype
         wire_protocol_client.status_blob.vm_status = vmstatus
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
 
         with patch.object(StatusBlob, "prepare", side_effect=Exception) as mock_prepare:
             self.assertRaises(ProtocolError, wire_protocol_client.upload_status_blob)
@@ -332,7 +336,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.artifacts_profile_blob = testurl
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
         wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
 
         with patch.object(HostPluginProtocol, "get_artifact_request",
@@ -361,7 +365,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.artifacts_profile_blob = testurl
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
         wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
 
         # response is invalid json
@@ -381,7 +385,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.artifacts_profile_blob = testurl
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
         wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
 
         wire_protocol_client.call_storage_service = Mock(
@@ -446,7 +450,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client = WireProtocol(wireserver_url).client
         wire_protocol_client.ext_conf = ExtensionsConfig(None)
         wire_protocol_client.ext_conf.artifacts_profile_blob = testurl
-        goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
+        goal_state = GoalState(mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state)
         wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
         wire_protocol_client.fetch = Mock(side_effect=[None, '{"onHold": "true"}'])
         with patch.object(HostPluginProtocol,
@@ -569,7 +573,7 @@ class TestWireClient(AgentTestCase):
         goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
         self.assertFalse(os.path.exists(goal_state_file))
 
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         client = WireClient(wireserver_url)
         client.save_or_update_goal_state_file(incarnation, xml_text)
 
@@ -582,7 +586,7 @@ class TestWireClient(AgentTestCase):
     def test_save_or_update_goal_state_should_update_existing_goal_state_file(self):
         incarnation = 42
         goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
 
         with open(goal_state_file, "w") as f:
             f.write(xml_text)
@@ -594,7 +598,7 @@ class TestWireClient(AgentTestCase):
             self.assertEquals("".join(contents), xml_text)
 
         # Update the container id
-        new_goal_state = WireProtocolData(DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
+        new_goal_state = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
                                                                         "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2")
         client = WireClient(wireserver_url)
         client.save_or_update_goal_state_file(incarnation, new_goal_state)
@@ -611,7 +615,7 @@ class TestWireClient(AgentTestCase):
         with open(incarnation_file, "w") as f:
             f.write(incarnation)
 
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         goal_state_file = os.path.join(conf.get_lib_dir(), "GoalState.{0}.xml".format(incarnation))
 
         with open(goal_state_file, "w") as f:
@@ -622,7 +626,7 @@ class TestWireClient(AgentTestCase):
         old_container_id = host.container_id
 
         # Update the container id
-        new_goal_state = WireProtocolData(DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
+        new_goal_state = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state.replace("c6d5526c-5ac2-4200-b6e2-56f2b70c5ab2",
                                                                         "z6d5526c-5ac2-4200-b6e2-56f2b70c5ab2")
         with patch("azurelinuxagent.common.protocol.wire.WireClient.fetch_config", return_value=new_goal_state):
             client.update_goal_state(forced=False)
@@ -975,7 +979,7 @@ class TestWireClient(AgentTestCase):
                         self.assertEquals(HostPluginProtocol.is_default_channel(), False)
 
     def test_send_request_using_appropriate_channel_should_not_invoke_host_channel_when_direct_channel_succeeds(self, *args):
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         client = WireClient(wireserver_url)
         client.goal_state = GoalState(xml_text)
         client.get_host_plugin().set_default_channel(False)
@@ -998,7 +1002,7 @@ class TestWireClient(AgentTestCase):
         self.assertEquals(0, host_func.counter)
 
     def test_send_request_using_appropriate_channel_should_not_use_direct_channel_when_host_channel_is_default(self, *args):
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         client = WireClient(wireserver_url)
         client.goal_state = GoalState(xml_text)
         client.get_host_plugin().set_default_channel(True)
@@ -1021,7 +1025,7 @@ class TestWireClient(AgentTestCase):
         self.assertEquals(1, host_func.counter)
 
     def test_send_request_using_appropriate_channel_should_use_host_channel_when_direct_channel_fails(self, *args):
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         client = WireClient(wireserver_url)
         client.goal_state = GoalState(xml_text)
         host = client.get_host_plugin()
@@ -1047,7 +1051,7 @@ class TestWireClient(AgentTestCase):
         self.assertEquals(True, host.is_default_channel())
 
     def test_send_request_using_appropriate_channel_should_retry_the_host_channel_after_reloading_goal_state(self, *args):
-        xml_text = WireProtocolData(DATA_FILE).goal_state
+        xml_text = mockwiredata.WireProtocolData(mockwiredata.DATA_FILE).goal_state
         client = WireClient(wireserver_url)
         client.goal_state = GoalState(xml_text)
         client.get_host_plugin().set_default_channel(False)
diff --git a/tests/test_agent.py b/tests/test_agent.py
index d0122641..dedc9905 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -20,7 +20,7 @@ import os.path
 from azurelinuxagent.agent import *
 from azurelinuxagent.common.conf import *
 
-from tests.tools import *
+from tests.tools import AgentTestCase, data_dir, Mock, patch
 
 EXPECTED_CONFIGURATION = \
 """AutoUpdate.Enabled = True
@@ -70,6 +70,7 @@ ResourceDisk.MountOptions = None
 ResourceDisk.MountPoint = /mnt/resource
 ResourceDisk.SwapSizeMB = 0""".split('\n')
 
+
 class TestAgent(AgentTestCase):
 
     def test_accepts_configuration_path(self):
diff --git a/tests/test_import.py b/tests/test_import.py
index c792c241..5d77d37c 100644
--- a/tests/test_import.py
+++ b/tests/test_import.py
@@ -1,4 +1,4 @@
-from tests.tools import *
+from tests.tools import AgentTestCase
 import azurelinuxagent.common.osutil as osutil
 import azurelinuxagent.common.dhcp as dhcp
 import azurelinuxagent.common.protocol as protocol
diff --git a/tests/utils/test_archive.py b/tests/utils/test_archive.py
index 1016c9aa..22949e48 100644
--- a/tests/utils/test_archive.py
+++ b/tests/utils/test_archive.py
@@ -3,9 +3,14 @@
 from datetime import datetime, timedelta
 
 import zipfile
+import os
+import shutil
+import tempfile
 
+import azurelinuxagent.common.logger as logger
+from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.archive import StateFlusher, StateArchiver, MAX_ARCHIVED_STATES
-from tests.tools import *
+from tests.tools import AgentTestCase
 
 debug = False
 if os.environ.get('DEBUG') == '1':
diff --git a/tests/utils/test_crypt_util.py b/tests/utils/test_crypt_util.py
index c9ae9bc0..c923c176 100644
--- a/tests/utils/test_crypt_util.py
+++ b/tests/utils/test_crypt_util.py
@@ -15,9 +15,14 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import os
+import sys
+import unittest
+
+import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.exception import CryptError
 from azurelinuxagent.common.utils.cryptutil import CryptUtil
-from tests.tools import *
+from tests.tools import AgentTestCase, data_dir, load_data, skip_if_predicate_true
 
 
 def is_python_version_26():
diff --git a/tests/utils/test_extension_process_util.py b/tests/utils/test_extension_process_util.py
index 2e036e74..c51f78af 100644
--- a/tests/utils/test_extension_process_util.py
+++ b/tests/utils/test_extension_process_util.py
@@ -18,8 +18,11 @@ from azurelinuxagent.common.exception import ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils.extensionprocessutil import format_stdout_stderr, read_output, \
     wait_for_process_completion_or_timeout, handle_process_completion
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
+import os
+import shutil
 import subprocess
+import tempfile
 
 
 class TestProcessUtils(AgentTestCase):
diff --git a/tests/utils/test_file_util.py b/tests/utils/test_file_util.py
index a3aadf2a..c3d4fdb9 100644
--- a/tests/utils/test_file_util.py
+++ b/tests/utils/test_file_util.py
@@ -21,11 +21,14 @@ import random
 import string
 import tempfile
 import uuid
+import os
+import shutil
+import unittest
 
 import azurelinuxagent.common.utils.fileutil as fileutil
 
 from azurelinuxagent.common.future import ustr
-from tests.tools import *
+from tests.tools import AgentTestCase, patch
 
 
 class TestFileOperations(AgentTestCase):
@@ -324,5 +327,6 @@ DHCP_HOSTNAME=test\n"
             self.assertFalse(os.path.isdir(d))
             self.assertFalse(os.path.isfile(d))
 
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/utils/test_network_util.py b/tests/utils/test_network_util.py
index 8dfefa22..ba5c630c 100644
--- a/tests/utils/test_network_util.py
+++ b/tests/utils/test_network_util.py
@@ -17,7 +17,7 @@
 
 import azurelinuxagent.common.utils.networkutil as networkutil
 
-from tests.tools import *
+from tests.tools import AgentTestCase
 
 
 class TestNetworkOperations(AgentTestCase):
diff --git a/tests/utils/test_rest_util.py b/tests/utils/test_rest_util.py
index f097f10f..97d1903b 100644
--- a/tests/utils/test_rest_util.py
+++ b/tests/utils/test_rest_util.py
@@ -15,11 +15,14 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
+import os
+import unittest
+
 from azurelinuxagent.common.exception import HttpError, ResourceGoneError, InvalidContainerError
 import azurelinuxagent.common.utils.restutil as restutil
 from azurelinuxagent.common.utils.restutil import HTTP_USER_AGENT
 from azurelinuxagent.common.future import httpclient, ustr
-from tests.tools import *
+from tests.tools import AgentTestCase, call, Mock, MagicMock, patch
 
 
 class TestIOErrorCounter(AgentTestCase):
@@ -28,7 +31,7 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
 
         counts = restutil.IOErrorCounter.get_and_reset()
         self.assertEqual(1, counts["hostplugin"])
@@ -40,7 +43,7 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+            restutil.KNOWN_WIRESERVER_IP, 80)
 
         counts = restutil.IOErrorCounter.get_and_reset()
         self.assertEqual(0, counts["hostplugin"])
@@ -64,11 +67,11 @@ class TestIOErrorCounter(AgentTestCase):
         restutil.IOErrorCounter.set_protocol_endpoint()
 
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+            restutil.KNOWN_WIRESERVER_IP, restutil.HOST_PLUGIN_PORT)
         restutil.IOErrorCounter.increment(
-            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+            restutil.KNOWN_WIRESERVER_IP, 80)
         restutil.IOErrorCounter.increment(
             '169.254.169.254', 80)
         restutil.IOErrorCounter.increment(
diff --git a/tests/utils/test_shell_util.py b/tests/utils/test_shell_util.py
index 145fb2d0..a8f545b5 100644
--- a/tests/utils/test_shell_util.py
+++ b/tests/utils/test_shell_util.py
@@ -16,8 +16,7 @@
 # Requires Python 2.6+ and Openssl 1.0+
 #
 
-from tests.tools import *
-from azurelinuxagent.common.logger import LogLevel
+from tests.tools import AgentTestCase, patch
 import unittest
 import azurelinuxagent.common.utils.shellutil as shellutil
 
@@ -155,7 +154,8 @@ class RunCommandTestCase(AgentTestCase):
         exception = context_manager.exception
         self.assertIn("No such file or directory", str(exception))
 
-    def test_run_command_it_should_not_log_by_default(self):
+    @patch("azurelinuxagent.common.utils.shellutil.logger", autospec=True)
+    def test_run_command_it_should_not_log_by_default(self, mock_logger):
 
         def assert_no_message_logged(command):
             try:
diff --git a/tests/utils/test_text_util.py b/tests/utils/test_text_util.py
index a32058e8..0701704a 100644
--- a/tests/utils/test_text_util.py
+++ b/tests/utils/test_text_util.py
@@ -16,9 +16,11 @@
 #
 
 from distutils.version import LooseVersion as Version
-from tests.tools import *
+from tests.tools import AgentTestCase
 
 import hashlib
+import os
+import unittest
 
 import azurelinuxagent.common.utils.textutil as textutil