Repository: Azure/WALinuxAgent
Pull request #437 (2016-09-17 00:13:49)

Replace unnecessary shell out logics with python based ones (#232)
Throughout the code, there exist several parts invoking shell commands to do certain tasks. As 'shelling out' is expensive and often unreliable, we need to refactor them to use python provided functions instead.

Unit tests are also added to validate updated/newly added codes are working as expected.


Patch:
diff --git a/azurelinuxagent/agent.py b/azurelinuxagent/agent.py
index 2c21e47f..553ac3d1 100644
--- a/azurelinuxagent/agent.py
+++ b/azurelinuxagent/agent.py
@@ -31,7 +31,7 @@ import azurelinuxagent.common.conf as conf
 from azurelinuxagent.common.version import AGENT_NAME, AGENT_LONG_VERSION, \
                                      DISTRO_NAME, DISTRO_VERSION, \
                                      PY_VERSION_MAJOR, PY_VERSION_MINOR, \
-                                     PY_VERSION_MICRO, GOAL_STATE_AGENT_VERSION
+                                     PY_VERSION_MICRO
 from azurelinuxagent.common.osutil import get_osutil
 
 class Agent(object):
@@ -173,14 +173,10 @@ def version():
     """
     Show agent version
     """
-    print(("{0} running on {1} {2}".format(AGENT_LONG_VERSION,
-                                           DISTRO_NAME,
-                                           DISTRO_VERSION)))
-    print("Python: {0}.{1}.{2}".format(PY_VERSION_MAJOR,
-                                       PY_VERSION_MINOR,
+    print(("{0} running on {1} {2}".format(AGENT_LONG_VERSION, DISTRO_NAME,
+                                          DISTRO_VERSION)))
+    print("Python: {0}.{1}.{2}".format(PY_VERSION_MAJOR, PY_VERSION_MINOR,
                                        PY_VERSION_MICRO))
-    print("Goal state agent: {0}".format(GOAL_STATE_AGENT_VERSION))
-
 def usage():
     """
     Show agent usage
diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py
index 3204ead3..73eaf23a 100644
--- a/azurelinuxagent/common/osutil/default.py
+++ b/azurelinuxagent/common/osutil/default.py
@@ -16,6 +16,7 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+import multiprocessing
 import os
 import re
 import shutil
@@ -247,6 +248,16 @@ class DefaultOSUtil(object):
         else:
             return False
 
+    def set_selinux_enforce(self, state):
+        """
+        Calls shell command 'setenforce' with 'state'
+        and returns resulting exit code.
+        """
+        if self.is_selinux_system():
+            if state: s = '1'
+            else: s='0'
+            return shellutil.run("setenforce "+s)
+
     def set_selinux_context(self, path, con):
         """
         Calls shell 'chcon' with 'path' and 'con' context.
@@ -751,19 +762,11 @@ class DefaultOSUtil(object):
         return base64.b64decode(data)
 
     def get_total_mem(self):
-        cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'"
-        ret = shellutil.run_get_output(cmd)
-        if ret[0] == 0:
-            return int(ret[1])/1024
-        else:
-            raise OSUtilError("Failed to get total memory: {0}".format(ret[1]))
+        # Get total memory in bytes and divide by 1024**2 to get the valu in MB.
+        return os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024**2)
 
     def get_processor_cores(self):
-        ret = shellutil.run_get_output("grep 'processor.*:' /proc/cpuinfo |wc -l")
-        if ret[0] == 0:
-            return int(ret[1])
-        else:
-            raise OSUtilError("Failed to get processor cores")
+        return multiprocessing.cpu_count()
 
     def set_admin_access_to_ip(self, dest_ip):
         #This allows root to access dest_ip
diff --git a/azurelinuxagent/common/osutil/freebsd.py b/azurelinuxagent/common/osutil/freebsd.py
index 54c74525..362aa92a 100644
--- a/azurelinuxagent/common/osutil/freebsd.py
+++ b/azurelinuxagent/common/osutil/freebsd.py
@@ -22,7 +22,7 @@ import azurelinuxagent.common.utils.textutil as textutil
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.exception import OSUtilError
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
-from azurelinuxagent.common.future import ustr
+
 
 class FreeBSDOSUtil(DefaultOSUtil):
     def __init__(self):
@@ -196,50 +196,3 @@ class FreeBSDOSUtil(DefaultOSUtil):
         logger.verbose("Interface info: ({0},{1},{2})", iface, inet, mac)
 
         return iface, inet, mac
-
-    def device_for_ide_port(self, port_id):
-        """
-        Return device name attached to ide port 'n'.
-        """
-        if port_id > 3:
-            return None
-        g0 = "00000000"
-        if port_id > 1:
-            g0 = "00000001"
-            port_id = port_id - 2
-        err, output = shellutil.run_get_output('sysctl dev.storvsc | grep pnpinfo | grep deviceid=')
-        if err:
-            return None
-        g1 = "000" + ustr(port_id)
-        g0g1 = "{0}-{1}".format(g0, g1)
-        """
-        search 'X' from 'dev.storvsc.X.%pnpinfo: classid=32412632-86cb-44a2-9b5c-50d1417354f5 deviceid=00000000-0001-8899-0000-000000000000'
-        """
-        cmd_search_ide = "sysctl dev.storvsc | grep pnpinfo | grep deviceid={0}".format(g0g1)
-        err, output = shellutil.run_get_output(cmd_search_ide)
-        if err:
-            return None
-        cmd_extract_id = cmd_search_ide + "|awk -F . '{print $3}'"
-        err, output = shellutil.run_get_output(cmd_extract_id)
-        """
-        try to search 'blkvscX' and 'storvscX' to find device name
-        """
-        output = output.rstrip()
-        cmd_search_blkvsc = "camcontrol devlist -b | grep blkvsc{0} | awk '{{print $1}}'".format(output)
-        err, output = shellutil.run_get_output(cmd_search_blkvsc)
-        if err == 0:
-            output = output.rstrip()
-            cmd_search_dev="camcontrol devlist | grep {0} | awk -F \( '{{print $2}}'|awk -F , '{{print $1}}'".format(output)
-            err, output = shellutil.run_get_output(cmd_search_dev)
-            if err == 0:
-                return output.rstrip()
-
-        cmd_search_storvsc = "camcontrol devlist -b | grep storvsc{0} | awk '{{print $1}}'".format(output)
-        err, output = shellutil.run_get_output(cmd_search_storvsc)
-        if err == 0:
-            output = output.rstrip()
-            cmd_search_dev="camcontrol devlist | grep {0} | awk -F \( '{{print $2}}'|awk -F , '{{print $1}}'".format(output)
-            err, output = shellutil.run_get_output(cmd_search_dev)
-            if err == 0:
-                return output.rstrip()
-        return None
diff --git a/azurelinuxagent/common/protocol/metadata.py b/azurelinuxagent/common/protocol/metadata.py
index 448649ae..071f1f7d 100644
--- a/azurelinuxagent/common/protocol/metadata.py
+++ b/azurelinuxagent/common/protocol/metadata.py
@@ -16,6 +16,7 @@
 #
 # Requires Python 2.4+ and Openssl 1.0+
 
+import base64
 import json
 import os
 import shutil
@@ -294,8 +295,11 @@ class Certificates(object):
         p7b_file = os.path.join(conf.get_lib_dir(), P7B_FILE_NAME)
 
         # Wrapping the certificate lines.
-        b64_cmd = "echo {0} | base64 -d > {1}"
-        shellutil.run(b64_cmd.format(data, p7b_file))
+        # decode and save the result into p7b_file
+        fileStream = open(p7b_file, 'w')
+        fileStream.write(textutil.b64decode(data))
+        fileStream.close()
+
         ssl_cmd = "openssl pkcs7 -text -in {0} -inform der | grep -v '^-----' "
         ret, data = shellutil.run_get_output(ssl_cmd.format(p7b_file))
 
diff --git a/azurelinuxagent/common/protocol/restapi.py b/azurelinuxagent/common/protocol/restapi.py
index aa7872bb..3a30be12 100644
--- a/azurelinuxagent/common/protocol/restapi.py
+++ b/azurelinuxagent/common/protocol/restapi.py
@@ -101,12 +101,11 @@ class CertificateData(DataContract):
         self.certificateData = certificateData
 
 class Cert(DataContract):
-    def __init__(self, name=None, thumbprint=None, certificateDataUri=None, storeName=None, storeLocation=None):
+    def __init__(self, name=None, thumbprint=None, certificateDataUri=None):
         self.name = name
         self.thumbprint = thumbprint
         self.certificateDataUri = certificateDataUri
-        self.storeLocation = storeLocation
-        self.storeName = storeName
+
 
 class CertList(DataContract):
     def __init__(self):
diff --git a/azurelinuxagent/common/utils/fileutil.py b/azurelinuxagent/common/utils/fileutil.py
index 7ef4fef1..b0b6fb75 100644
--- a/azurelinuxagent/common/utils/fileutil.py
+++ b/azurelinuxagent/common/utils/fileutil.py
@@ -21,11 +21,11 @@
 File operation util functions
 """
 
+import glob
 import os
 import re
 import shutil
 import pwd
-import tempfile
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.future import ustr
 import azurelinuxagent.common.utils.textutil as textutil
@@ -111,9 +111,11 @@ def chmod(path, mode):
         os.chmod(path, mode)
 
 def rm_files(*args):
-    for path in args:
-        if os.path.isfile(path):
-            os.remove(path)
+    for paths in args:
+        #Find all possible file paths
+        for path in glob.glob(paths):
+            if os.path.isfile(path):
+                os.remove(path)
 
 def rm_dirs(*args):
     """
@@ -169,3 +171,12 @@ def findstr_in_file(file_path, pattern_str):
 
     return None
 
+def get_all_files(root_path):
+    """
+    Find all files under the given root path
+    """
+    result = []
+    for root, dirs, files in os.walk(root_path):
+        result.extend([os.path.join(root, file) for file in files])
+
+    return result
diff --git a/azurelinuxagent/common/utils/textutil.py b/azurelinuxagent/common/utils/textutil.py
index 6d460cb1..db320a62 100644
--- a/azurelinuxagent/common/utils/textutil.py
+++ b/azurelinuxagent/common/utils/textutil.py
@@ -278,6 +278,11 @@ def b64encode(s):
         return base64.b64encode(bytes(s, 'utf-8')).decode('utf-8')
     return base64.b64encode(s)
 
+def b64decode(s):
+    from azurelinuxagent.common.version import PY_VERSION_MAJOR
+    if PY_VERSION_MAJOR > 2:
+        return base64.b64decode(s).decode('utf-8')
+    return base64.b64decode(s)
 
 def safe_shlex_split(s):
     import shlex
diff --git a/azurelinuxagent/common/version.py b/azurelinuxagent/common/version.py
index 7d19e4fa..0ba078eb 100644
--- a/azurelinuxagent/common/version.py
+++ b/azurelinuxagent/common/version.py
@@ -51,7 +51,7 @@ def get_distro():
 
 AGENT_NAME = "WALinuxAgent"
 AGENT_LONG_NAME = "Azure Linux Agent"
-AGENT_VERSION = '2.1.6.5'
+AGENT_VERSION = '2.1.6.1'
 AGENT_LONG_VERSION = "{0}-{1}".format(AGENT_NAME, AGENT_VERSION)
 AGENT_DESCRIPTION = """\
 The Azure Linux Agent supports the provisioning and running of Linux
@@ -66,8 +66,6 @@ AGENT_PATTERN = "{0}-(.*)".format(AGENT_NAME)
 AGENT_NAME_PATTERN = re.compile(AGENT_PATTERN)
 AGENT_DIR_PATTERN = re.compile(".*/{0}".format(AGENT_PATTERN))
 
-EXT_HANDLER_PATTERN = b".*/WALinuxAgent-(\w.\w.\w[.\w]*)-.*-run-exthandlers"
-EXT_HANDLER_REGEX = re.compile(EXT_HANDLER_PATTERN)
 
 # Set the CURRENT_AGENT and CURRENT_VERSION to match the agent directory name
 # - This ensures the agent will "see itself" using the same name and version
@@ -86,23 +84,6 @@ def set_current_agent():
     return agent, FlexibleVersion(version)
 CURRENT_AGENT, CURRENT_VERSION = set_current_agent()
 
-def set_goal_state_agent():
-    agent = None
-    pids = [pid for pid in os.listdir('/proc') if pid.isdigit()]
-    for pid in pids:
-        try:
-            pname = open(os.path.join('/proc', pid, 'cmdline'), 'rb').read()
-            match = EXT_HANDLER_REGEX.match(pname)
-            if match:
-                agent = match.group(1)
-                break
-        except IOError:
-            continue
-    if agent is None:
-        agent = CURRENT_VERSION
-    return agent
-GOAL_STATE_AGENT_VERSION = set_goal_state_agent()
-
 def is_current_agent_installed():
     return CURRENT_AGENT == AGENT_LONG_VERSION
 
diff --git a/azurelinuxagent/daemon/resourcedisk/freebsd.py b/azurelinuxagent/daemon/resourcedisk/freebsd.py
index 4ca00586..36a3ac98 100644
--- a/azurelinuxagent/daemon/resourcedisk/freebsd.py
+++ b/azurelinuxagent/daemon/resourcedisk/freebsd.py
@@ -57,24 +57,19 @@ class FreeBSDResourceDiskHandler(ResourceDiskHandler):
             raise ResourceDiskError("Unable to detect resource disk device:{0}".format(output))
         disks = self.parse_gpart_list(output)
 
-        device = self.osutil.device_for_ide_port(1)
-        if device is None:
-        # fallback logic to find device
-            err, output = shellutil.run_get_output('camcontrol periphlist 2:1:0')
-            if err:
-                # try again on "3:1:0"
-                err, output = shellutil.run_get_output('camcontrol periphlist 3:1:0')
-                if err:
-                    raise ResourceDiskError("Unable to detect resource disk device:{0}".format(output))
+        err, output = shellutil.run_get_output('camcontrol periphlist 2:1:0')
+        if err:
+            raise ResourceDiskError("Unable to detect resource disk device:{0}".format(output))
 
         # 'da1:  generation: 4 index: 1 status: MORE\npass2:  generation: 4 index: 2 status: LAST\n'
-            for line in output.split('\n'):
-                index = line.find(':')
-                if index > 0:
-                    geom_name = line[:index]
-                    if geom_name in disks:
-                        device = geom_name
-                        break
+        device = None
+        for line in output.split('\n'):
+            index = line.find(':')
+            if index > 0:
+                geom_name = line[:index]
+                if geom_name in disks:
+                    device = geom_name
+                    break
 
         if not device:
             raise ResourceDiskError("Unable to detect resource disk device.")
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 551ac731..5f75dc5b 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -21,6 +21,7 @@ import glob
 import json
 import os
 import shutil
+import stat
 import subprocess
 import time
 import zipfile
@@ -174,12 +175,16 @@ class ExtHandlersHandler(object):
             add_event(AGENT_NAME, version=CURRENT_VERSION, is_success=False, message=msg)
             return
 
-        msg = u"Handle extensions updates for incarnation {0}".format(etag)
-        logger.verbose(msg)
-        # Log status report success on new config
-        self.log_report = True
-        self.handle_ext_handlers(etag)
-        self.last_etag = etag
+        if self.last_etag is not None and self.last_etag == etag:
+            msg = u"Incarnation {0} has no extension updates".format(etag)
+            logger.verbose(msg)
+            self.log_report = False
+        else:
+            msg = u"Handle extensions updates for incarnation {0}".format(etag)
+            logger.info(msg)
+            self.log_report = True #Log status report success on new config
+            self.handle_ext_handlers()
+            self.last_etag = etag
 
         self.report_ext_handlers_status()
 
@@ -187,24 +192,18 @@ class ExtHandlersHandler(object):
         self.report_ext_handlers_status()
         return
    
-    def handle_ext_handlers(self, etag=None):
+    def handle_ext_handlers(self):
         if self.ext_handlers.extHandlers is None or \
                 len(self.ext_handlers.extHandlers) == 0:
-            logger.verbose("No ext handler config found")
+            logger.info("No ext handler config found")
             return
 
         for ext_handler in self.ext_handlers.extHandlers:
             #TODO handle install in sequence, enable in parallel
-            self.handle_ext_handler(ext_handler, etag)
+            self.handle_ext_handler(ext_handler)
     
-    def handle_ext_handler(self, ext_handler, etag):
+    def handle_ext_handler(self, ext_handler):
         ext_handler_i = ExtHandlerInstance(ext_handler, self.protocol)
-
-        ext_handler_i.decide_version()
-        if not ext_handler_i.is_upgrade and self.last_etag == etag:
-            ext_handler_i.logger.info("No upgrade detected for etag {0}", etag)
-            return
-
         try:
             state = ext_handler.properties.state
             ext_handler_i.logger.info("Expected handler state: {0}", state)
@@ -223,10 +222,12 @@ class ExtHandlersHandler(object):
     
     def handle_enable(self, ext_handler_i):
 
+        ext_handler_i.decide_version() 
+
         old_ext_handler_i = ext_handler_i.get_installed_ext_handler()
         if old_ext_handler_i is not None and \
            old_ext_handler_i.version_gt(ext_handler_i):
-            raise ExtensionError(u"Downgrade not allowed")
+            raise ExtensionError(u"Downgrade not allowed")  
 
         handler_state = ext_handler_i.get_handler_state()
         ext_handler_i.logger.info("Current handler state is: {0}", handler_state)
@@ -326,7 +327,6 @@ class ExtHandlerInstance(object):
         self.protocol = protocol
         self.operation = None
         self.pkg = None
-        self.is_upgrade = False
 
         prefix = "[{0}]".format(self.get_full_name())
         self.logger = logger.Logger(logger.DEFAULT_LOGGER, prefix)
@@ -432,11 +432,6 @@ class ExtHandlerInstance(object):
             self.pkg = selected_pkg
             self.ext_handler.properties.version = selected_pkg.version
 
-        # Note if the selected package is greater than that installed
-        if installed_pkg is None \
-            or FlexibleVersion(self.pkg.version) > FlexibleVersion(installed_pkg.version):
-            self.is_upgrade = True
-
         if self.pkg is None:
             raise ExtensionError("Failed to find any valid extension package")
 
@@ -524,8 +519,10 @@ class ExtHandlerInstance(object):
         except IOError as e:
             raise ExtensionError(u"Failed to write and unzip plugin", e)
 
-        chmod = "find {0} -type f | xargs chmod u+x".format(self.get_base_dir())
-        shellutil.run(chmod)
+        #Add user execute permission to all files under the base dir
+        for file in fileutil.get_all_files(self.get_base_dir()):
+            fileutil.chmod(file, os.stat(file).st_mode | stat.S_IXUSR)
+
         self.report_event(message="Download succeeded")
 
         self.logger.info("Initialize extension directory")
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index be481950..8cc0f026 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -85,6 +85,7 @@ class UpdateHandler(object):
         self.protocol_util = get_protocol_util()
 
         self.running = True
+        self.last_etag = None
         self.last_attempt_time = None
 
         self.agents = []
@@ -345,6 +346,10 @@ class UpdateHandler(object):
                 message=msg)
             return False
 
+        if self.last_etag is not None and self.last_etag == etag:
+            logger.info(u"Incarnation {0} has no agent updates", etag)
+            return False
+
         manifests = [m for m in manifest_list.vmAgentManifests if m.family == family]
         if len(manifests) == 0:
             logger.info(u"Incarnation {0} has no agent family {1} updates", etag, family)
diff --git a/azurelinuxagent/pa/deprovision/default.py b/azurelinuxagent/pa/deprovision/default.py
index 3a916e21..3d60c454 100644
--- a/azurelinuxagent/pa/deprovision/default.py
+++ b/azurelinuxagent/pa/deprovision/default.py
@@ -63,8 +63,8 @@ class DeprovisionHandler(object):
 
     def regen_ssh_host_key(self, warnings, actions):
         warnings.append("WARNING! All SSH host key pairs will be deleted.")
-        actions.append(DeprovisionAction(shellutil.run,
-                                         ['rm -f /etc/ssh/ssh_host_*key*']))
+        actions.append(DeprovisionAction(fileutil.rm_files,
+                                         ['/etc/ssh/ssh_host_*key*']))
 
     def stop_agent_service(self, warnings, actions):
         warnings.append("WARNING! The waagent service will be stopped.")
diff --git a/azurelinuxagent/pa/provision/default.py b/azurelinuxagent/pa/provision/default.py
index 454a54c2..c7514c37 100644
--- a/azurelinuxagent/pa/provision/default.py
+++ b/azurelinuxagent/pa/provision/default.py
@@ -81,7 +81,7 @@ class ProvisionHandler(object):
     def reg_ssh_host_key(self):
         keypair_type = conf.get_ssh_host_keypair_type()
         if conf.get_regenerate_ssh_host_key():
-            shellutil.run("rm -f /etc/ssh/ssh_host_*key*")
+            fileutil.rm_files("/etc/ssh/ssh_host_*key*")
             keygen_cmd = "ssh-keygen -N '' -t {0} -f /etc/ssh/ssh_host_{1}_key"
             shellutil.run(keygen_cmd.format(keypair_type, keypair_type))
         thumbprint = self.get_ssh_host_key_thumbprint(keypair_type)
@@ -134,9 +134,17 @@ class ProvisionHandler(object):
         logger.info("Configure sshd")
         self.osutil.conf_sshd(ovfenv.disable_ssh_password_auth)
 
+        # Disable selinux temporary
+        sel = self.osutil.is_selinux_enforcing()
+        if sel:
+            self.osutil.set_selinux_enforce(0)
+
         self.deploy_ssh_pubkeys(ovfenv)
         self.deploy_ssh_keypairs(ovfenv)
 
+        if sel:
+            self.osutil.set_selinux_enforce(1)
+
     def save_customdata(self, ovfenv):
         customdata = ovfenv.customdata
         if customdata is None:
diff --git a/init/ubuntu/walinuxagent.service b/init/ubuntu/walinuxagent.service
index 3cf47c5b..ae8dccb5 100755
--- a/init/ubuntu/walinuxagent.service
+++ b/init/ubuntu/walinuxagent.service
@@ -7,8 +7,8 @@
 [Unit]
 Description=Azure Linux Agent
 
-After=network-online.target cloud-init.service
-Wants=network-online.target sshd.service sshd-keygen.service
+After=network-online.target cloud-final.service
+Wants=network-online.target sshd.service sshd-keygen.service cloud-final.service
 
 ConditionFileIsExecutable=/usr/sbin/waagent
 ConditionPathExists=/etc/waagent.conf


Test patch:
diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py
index d9d00f64..d982b7e2 100644
--- a/tests/common/osutil/test_default.py
+++ b/tests/common/osutil/test_default.py
@@ -142,5 +142,27 @@ class TestOSUtil(AgentTestCase):
                 self.assertTrue(endpoint is not None)
                 self.assertEqual(endpoint, "second")
 
+    def test_get_total_mem(self):
+        """
+        Validate the returned value matches to the one retrieved by invoking shell command
+        """
+        cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'"
+        ret = shellutil.run_get_output(cmd)
+        if ret[0] == 0:
+            self.assertEqual(int(ret[1]) / 1024, get_osutil().get_total_mem())
+        else:
+            self.fail("Cannot retrieve total memory using shell command.")
+
+    def test_get_processor_cores(self):
+        """
+        Validate the returned value matches to the one retrieved by invoking shell command
+        """
+        cmd = "grep 'processor.*:' /proc/cpuinfo |wc -l"
+        ret = shellutil.run_get_output(cmd)
+        if ret[0] == 0:
+            self.assertEqual(int(ret[1]), get_osutil().get_processor_cores())
+        else:
+            self.fail("Cannot retrieve number of process cores using shell command.")
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 9e246b71..fe30ca15 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -162,7 +162,7 @@ class TestExtension(AgentTestCase):
         self.assertNotEquals(0, len(vm_status.vmAgent.extensionHandlers))
         handler_status = vm_status.vmAgent.extensionHandlers[0]
         self.assertEquals(expected_status, handler_status.status)
-        self.assertEquals("OSTCExtensions.ExampleHandlerLinux", 
+        self.assertEquals("OSTCExtensions.ExampleHandlerLinux",
                           handler_status.name)
         self.assertEquals(version, handler_status.version)
         self.assertEquals(expected_ext_count, len(handler_status.extensions))
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index b110b75c..1c3f002a 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -624,6 +624,7 @@ class TestUpdate(UpdateTestCase):
     def test_creation(self):
         self.assertTrue(self.update_handler.running)
 
+        self.assertEqual(None, self.update_handler.last_etag)
         self.assertEqual(None, self.update_handler.last_attempt_time)
 
         self.assertEqual(0, len(self.update_handler.agents))
@@ -664,14 +665,16 @@ class TestUpdate(UpdateTestCase):
         if versions is None or len(versions) <= 0:
             versions = [latest_version]
 
+        etag = self.update_handler.last_etag if self.update_handler.last_etag is not None else 42
         if protocol is None:
-            protocol = ProtocolMock(versions=versions)
+            protocol = ProtocolMock(etag=etag, versions=versions)
         self.update_handler.protocol_util = protocol
         conf.get_autoupdate_gafamily = Mock(return_value=protocol.family)
 
         return self.update_handler._upgrade_available(base_version=base_version)
 
     def test_ensure_latest_agent_returns_true_on_first_use(self):
+        self.assertEqual(None, self.update_handler.last_etag)
         self.assertTrue(self._test_upgrade_available())
         return
 
@@ -713,6 +716,11 @@ class TestUpdate(UpdateTestCase):
         self.assertFalse(self._test_upgrade_available())
         return
 
+    def test_ensure_latest_agent_skips_when_etag_matches(self):
+        self.update_handler.last_etag = 42
+        self.assertFalse(self._test_upgrade_available())
+        return
+
     def test_ensure_latest_agent_skips_if_when_no_new_versions(self):
         self.prepare_agents()
         base_version = self.agent_versions()[0] + 1
diff --git a/tests/tools.py b/tests/tools.py
index 2732f761..8801a0cd 100644
--- a/tests/tools.py
+++ b/tests/tools.py
@@ -56,6 +56,7 @@ class AgentTestCase(unittest.TestCase):
     def setUp(self):
         prefix = "{0}_".format(self.__class__.__name__)
         self.tmp_dir = tempfile.mkdtemp(prefix=prefix)
+        self.test_file = 'test_file'
         conf.get_autoupdate_enabled = Mock(return_value=True)
         conf.get_lib_dir = Mock(return_value=self.tmp_dir)
         ext_log_dir = os.path.join(self.tmp_dir, "azure")
diff --git a/tests/utils/test_file_util.py b/tests/utils/test_file_util.py
index 9a5479e2..f16f4099 100644
--- a/tests/utils/test_file_util.py
+++ b/tests/utils/test_file_util.py
@@ -24,8 +24,9 @@ from azurelinuxagent.common.future import ustr
 import azurelinuxagent.common.utils.fileutil as fileutil
 
 class TestFileOperations(AgentTestCase):
+
     def test_read_write_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = ustr(uuid.uuid4())
         fileutil.write_file(test_file, content)
 
@@ -34,7 +35,7 @@ class TestFileOperations(AgentTestCase):
         os.remove(test_file)
     
     def test_rw_utf8_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = u"\u6211"
         fileutil.write_file(test_file, content, encoding="utf-8")
 
@@ -43,14 +44,14 @@ class TestFileOperations(AgentTestCase):
         os.remove(test_file)
 
     def test_remove_bom(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         data = b'\xef\xbb\xbfhehe'
         fileutil.write_file(test_file, data, asbin=True)
         data = fileutil.read_file(test_file, remove_bom=True)
         self.assertNotEquals(0xbb, ord(data[0]))
    
     def test_append_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = ustr(uuid.uuid4())
         fileutil.append_file(test_file, content)
 
@@ -68,5 +69,53 @@ class TestFileOperations(AgentTestCase):
         filename = fileutil.base_name(filepath)
         self.assertEquals('abc', filename)
 
+    def test_remove_files(self):
+        import random
+        import string
+        import glob
+        random_word = lambda : ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
+
+        #Create 10 test files
+        test_file = os.path.join(self.tmp_dir, self.test_file)
+        test_file2 = os.path.join(self.tmp_dir, 'another_file')
+        test_files = [test_file + random_word() for _ in range(5)] + \
+                     [test_file2 + random_word() for _ in range(5)]
+        for file in test_files:
+            open(file, 'a').close()
+
+        #Remove files using fileutil.rm_files
+        test_file_pattern = test_file + '*'
+        test_file_pattern2 = test_file2 + '*'
+        fileutil.rm_files(test_file_pattern, test_file_pattern2)
+
+        self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern))))
+        self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern2))))
+
+    def test_get_all_files(self):
+        import random
+        import string
+        random_word = lambda: ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
+
+        # Create 10 test files at the root dir and 10 other in the sub dir
+        test_file = os.path.join(self.tmp_dir, self.test_file)
+        test_file2 = os.path.join(self.tmp_dir, 'another_file')
+        expected_files = [test_file + random_word() for _ in range(5)] + \
+                     [test_file2 + random_word() for _ in range(5)]
+
+        test_subdir = os.path.join(self.tmp_dir, 'test_dir')
+        os.mkdir(test_subdir)
+        test_file_in_subdir = os.path.join(test_subdir, self.test_file)
+        test_file_in_subdir2 = os.path.join(test_subdir, 'another_file')
+        expected_files.extend([test_file_in_subdir + random_word() for _ in range(5)] + \
+                     [test_file_in_subdir2 + random_word() for _ in range(5)])
+
+        for file in expected_files:
+            open(file, 'a').close()
+
+        # Get All files using fileutil.get_all_files
+        actual_files = fileutil.get_all_files(self.tmp_dir)
+
+        self.assertEqual(set(expected_files), set(actual_files))
+
 if __name__ == '__main__':
     unittest.main()