Repository: Azure/WALinuxAgent
Pull request #2160 (2021-02-05 21:42:53)

Set CPU quota on agent

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160?src=pr&el=h1) Report
> Merging [#2160](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160?src=pr&el=desc) (9d57c4e) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/e101f2082bd5621a01a448904eea20131b37f85d?el=desc) (e101f20) will **decrease** coverage by `0.01%`.
> The diff coverage is `72.72%`.

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

```diff
@@             Coverage Diff             @@
##           develop    #2160      +/-   ##
===========================================
- Coverage    70.65%   70.63%   -0.02%     
===========================================
  Files           96       96              
  Lines        13351    13357       +6     
  Branches      1882     1884       +2     
===========================================
+ Hits          9433     9435       +2     
- Misses        3504     3506       +2     
- Partials       414      416       +2     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `66.88% <72.72%> (+0.97%)` | :arrow_up: |
| [azurelinuxagent/ga/collect\_telemetry\_events.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160/diff?src=pr&el=tree#diff-YXp1cmVsaW51eGFnZW50L2dhL2NvbGxlY3RfdGVsZW1ldHJ5X2V2ZW50cy5weQ==) | `89.65% <0.00%> (-1.73%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160?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/2160?src=pr&el=footer). Last update [e101f20...9d57c4e](https://codecov.io/gh/Azure/WALinuxAgent/pull/2160?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 28695145..06c3e6ea 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -20,7 +20,7 @@ import re
 import subprocess
 
 from azurelinuxagent.common import logger
-from azurelinuxagent.common.cgroup import CpuCgroup, MemoryCgroup
+from azurelinuxagent.common.cgroup import CpuCgroup
 from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SystemdRunError
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.exception import ExtensionErrorCodes, CGroupsException
@@ -61,6 +61,13 @@ _AGENT_DROP_IN_FILE_CPU_ACCOUNTING_CONTENTS = """
 [Service]
 CPUAccounting=yes
 """
+_AGENT_DROP_IN_FILE_CPU_QUOTA = "12-CPUQuota.conf"
+_AGENT_DROP_IN_FILE_CPU_QUOTA_CONTENTS = """
+# This drop-in unit file was created by the Azure VM Agent.
+# Do not edit.
+[Service]
+CPUQuota=5%
+"""
 
 class CGroupConfigurator(object):
     """
@@ -106,22 +113,20 @@ class CGroupConfigurator(object):
                 if not self.__check_no_legacy_cgroups():
                     return
 
-                cpu_controller_root, memory_controller_root = self.__get_cgroup_controllers()
+                agent_unit_name = self._cgroups_api.get_agent_unit_name()
+                agent_slice = self._cgroups_api.get_unit_property(agent_unit_name, "Slice")
+                if agent_slice not in (_AZURE_SLICE, "system.slice"):
+                    self.__log_cgroup_warning("The agent is within an unexpected slice: {0}", agent_slice)
+                    return
 
-                agent_slice = self.__ensure_azure_slices_exist()
+                self.__setup_unit_files()
 
+                cpu_controller_root, memory_controller_root = self.__get_cgroup_controllers()
                 self._agent_cpu_cgroup_path, self._agent_memory_cgroup_path = self.__get_agent_cgroups(agent_slice, cpu_controller_root, memory_controller_root)
 
-                agent_service_name = self._cgroups_api.get_agent_unit_name()
                 if self._agent_cpu_cgroup_path is not None:
                     self.__log_cgroup_info("Agent CPU cgroup: {0}", self._agent_cpu_cgroup_path)
-                    CGroupsTelemetry.track_cgroup(CpuCgroup(agent_service_name, self._agent_cpu_cgroup_path))
-
-                if self._agent_memory_cgroup_path is not None:
-                    self.__log_cgroup_info("Agent Memory cgroup: {0}", self._agent_memory_cgroup_path)
-                    CGroupsTelemetry.track_cgroup(MemoryCgroup(agent_service_name, self._agent_memory_cgroup_path))
-
-                if self._agent_cpu_cgroup_path is not None or self._agent_memory_cgroup_path is not None:
+                    CGroupsTelemetry.track_cgroup(CpuCgroup(agent_unit_name, self._agent_cpu_cgroup_path))
                     self._cgroups_enabled = True
 
                 self.__log_cgroup_info('Cgroups enabled: {0}', self._cgroups_enabled)
@@ -235,7 +240,7 @@ class CGroupConfigurator(object):
 
             return cpu_controller_root, memory_controller_root
 
-        def __ensure_azure_slices_exist(self):
+        def __setup_unit_files(self):
             """
             The agent creates "azure.slice" for use by extensions and the agent. The agent runs under "azure.slice" directly and each
             extension runs under its own slice ("Microsoft.CPlat.Extension.slice" in the example below). All the slices for
@@ -254,8 +259,10 @@ class CGroupConfigurator(object):
 
             This method ensures that the "azure" and "vmextensions" slices are created. Setup should create those slices
             under /lib/systemd/system; but if they do not exist, __ensure_azure_slices_exist will create them.
-            It also creates drop-in files to set the agent's slice and CPU accounting properties if they have not been
+
+            It also creates drop-in files to set the agent's Slice, CPUAccounting, and CPUQuota properties if they have not been
             set up in the agent's unit file.
+
             Lastly, the method also cleans up unit files left over from previous versions of the agent.
 
             Returns the slice under which the agent should currently be running.
@@ -269,36 +276,42 @@ class CGroupConfigurator(object):
             agent_unit_file = "/lib/systemd/system/{0}".format(self._cgroups_api.get_agent_unit_name())
             agent_drop_in_file_slice = "/lib/systemd/system/{0}.d/{1}".format(self._cgroups_api.get_agent_unit_name(), _AGENT_DROP_IN_FILE_SLICE)
             agent_drop_in_file_cpu_accounting = "/lib/systemd/system/{0}.d/{1}".format(self._cgroups_api.get_agent_unit_name(), _AGENT_DROP_IN_FILE_CPU_ACCOUNTING)
+            agent_drop_in_file_cpu_quota = "/lib/systemd/system/{0}.d/{1}".format(self._cgroups_api.get_agent_unit_name(), _AGENT_DROP_IN_FILE_CPU_QUOTA)
+            agent_unit_name = self._cgroups_api.get_agent_unit_name()
 
             files_to_create = []
-            return_azure_slice = False
 
-            if os.path.exists(azure_slice):
-                return_azure_slice = True
-            else:
+            if not os.path.exists(azure_slice):
                 files_to_create.append((azure_slice, _AZURE_SLICE_CONTENTS))
 
             if not os.path.exists(vmextensions_slice):
                 files_to_create.append((vmextensions_slice, _VMEXTENSIONS_SLICE_CONTENTS))
 
-            if fileutil.findstr_in_file(agent_unit_file, "Slice=azure.slice"):
+            if fileutil.findre_in_file(agent_unit_file, r"Slice=") is not None:
                 self._cleanup_unit_file(agent_drop_in_file_slice)
             else:
                 if not os.path.exists(agent_drop_in_file_slice):
                     files_to_create.append((agent_drop_in_file_slice, _AGENT_DROP_IN_FILE_SLICE_CONTENTS))
 
-            if fileutil.findstr_in_file(agent_unit_file, "CPUAccounting=yes"):
+            if fileutil.findre_in_file(agent_unit_file, r"CPUAccounting=") is not None:
                 self._cleanup_unit_file(agent_drop_in_file_cpu_accounting)
             else:
                 if not os.path.exists(agent_drop_in_file_cpu_accounting):
                     files_to_create.append((agent_drop_in_file_cpu_accounting, _AGENT_DROP_IN_FILE_CPU_ACCOUNTING_CONTENTS))
 
+            if fileutil.findre_in_file(agent_unit_file, r"CPUQuota=") is not None:
+                self._cleanup_unit_file(agent_drop_in_file_cpu_quota)
+            else:
+                if not os.path.exists(agent_drop_in_file_cpu_quota):
+                    files_to_create.append((agent_drop_in_file_cpu_quota, _AGENT_DROP_IN_FILE_CPU_QUOTA_CONTENTS))
+
             if len(files_to_create) > 0:
                 try:
                     for path, contents in files_to_create:
                         self._create_unit_file(path, contents)
 
-                    # reload the systemd configuration, but the new slice will not be used until the agent's service restarts
+                    # reload the systemd configuration, but the new slice will not be used until the agent's service restarts,
+                    # but the changes in CPU quota will apply after the reload
                     try:
                         logger.info("Executing systemctl daemon-reload...")
                         shellutil.run_command(["systemctl", "daemon-reload"])
@@ -309,7 +322,7 @@ class CGroupConfigurator(object):
                     for unit_file in files_to_create:
                         self._cleanup_unit_file(unit_file)
 
-            return _AZURE_SLICE if return_azure_slice else "system.slice"
+            logger.info("Agent's CPUQuota: {0}",   self._cgroups_api.get_unit_property(agent_unit_name, "CPUQuotaPerSecUSec"))
 
         def _create_unit_file(self, path, contents):
             try:


Test patch:
diff --git a/tests/common/mock_cgroup_commands.py b/tests/common/mock_cgroup_commands.py
index 8bef9875..d999df9f 100644
--- a/tests/common/mock_cgroup_commands.py
+++ b/tests/common/mock_cgroup_commands.py
@@ -18,15 +18,15 @@
 import contextlib
 import os
 import re
+import shutil
 import subprocess
 import sys
 
-from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.utils import fileutil
 from tests.tools import patch, data_dir
 
-__MOCKED_COMMANDS = [
-    (r"^systemctl --version$",
+_MOCKED_COMMANDS = [
+   (r"^systemctl --version$",
 '''systemd 237
 +PAM +AUDIT +SELINUX +IMA +APPARMOR +SMACK +SYSVINIT +UTMP +LIBCRYPTSETUP +GCRYPT +GNUTLS +ACL +XZ +LZ4 +SECCOMP +BLKID +ELFUTILS +KMOD -IDN2 +IDN -PCRE2 default-hierarchy=hybrid
 '''),
@@ -48,11 +48,15 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
     (r"^mount -t cgroup2$",
 '''cgroup on /sys/fs/cgroup/unified type cgroup2 (rw,nosuid,nodev,noexec,relatime) 
 '''),
-
+    (r"^systemctl show walinuxagent\.service --property Slice",
+'''Slice=system.slice
+'''),
     (r"^systemctl show walinuxagent\.service --property CPUAccounting$",
 '''CPUAccounting=no
 '''),
-
+    (r"^systemctl show walinuxagent\.service --property CPUQuotaPerSecUSec$",
+'''CPUQuotaPerSecUSec=infinity
+'''),
     (r"^systemctl show walinuxagent\.service --property MemoryAccounting$",
 '''MemoryAccounting=no
 '''),
@@ -67,23 +71,28 @@ Thu 28 May 2020 07:25:55 AM PDT
 
 ]
 
-__service_file = "{0}.service".format(get_osutil().get_service_name())
-
-__MOCKED_FILES = [
-    (r"^/proc/self/cgroup$", os.path.join(data_dir, 'cgroups', 'proc_self_cgroup')),
-    (r"^/proc/[0-9]+/cgroup$", os.path.join(data_dir, 'cgroups', 'proc_pid_cgroup')),
-    (r"^/sys/fs/cgroup/unified/cgroup.controllers$", os.path.join(data_dir, 'cgroups', 'sys_fs_cgroup_unified_cgroup.controllers')),
-    (r"^/lib/systemd/system/{0}$".format(__service_file), os.path.join(data_dir, 'init', __service_file + "_system-slice")),
+_MOCKED_FILES = [
+    ("/proc/self/cgroup", os.path.join(data_dir, 'cgroups', 'proc_self_cgroup')),
+    (r"/proc/[0-9]+/cgroup", os.path.join(data_dir, 'cgroups', 'proc_pid_cgroup')),
+    ("/sys/fs/cgroup/unified/cgroup.controllers", os.path.join(data_dir, 'cgroups', 'sys_fs_cgroup_unified_cgroup.controllers'))
 ]
 
-__MOCKED_PATHS = [
+_MOCKED_PATHS = [
     r"^(/lib/systemd/system)",
     r"^(/etc/systemd/system)"
 ]
 
-# W0212: Access to a protected member __commands of a client class (protected-access) - Disabled: patcher.__commands is added
-# only for debugging purposes and should not be public (hence it is marked as private).
-# pylint: disable=protected-access
+
+#  Class has no __init__ method (no-init): Disabled; the class is just an enum
+class UnitFilePaths:  # pylint: disable=no-init
+    walinuxagent = "/lib/systemd/system/walinuxagent.service"
+    azure = "/lib/systemd/system/azure.slice"
+    vmextensions = "/lib/systemd/system/azure-vmextensions.slice"
+    slice = "/lib/systemd/system/walinuxagent.service.d/10-Slice.conf"
+    cpu_accounting = "/lib/systemd/system/walinuxagent.service.d/11-CPUAccounting.conf"
+    cpu_quota = "/lib/systemd/system/walinuxagent.service.d/12-CPUQuota.conf"
+
+
 @contextlib.contextmanager
 def mock_cgroup_commands(tmp_dir):
     """
@@ -103,15 +112,28 @@ def mock_cgroup_commands(tmp_dir):
     Matches are done using regular expressions; the regular expressions in __MOCKED_PATHS must create group 0 to indicate
     the section of the path that needs to be mapped (i.e. use parenthesis around the section that needs to be mapped.)
 
-    The command output used in __MOCKED_COMMANDS come from an Ubuntu 18 system.
+    The command output used in __MOCKED_COMMANDS comes from an Ubuntu 18 system.
+
+    The object returned by this function includes 3 methods (add_command_mock, add_path_mock, and add_file_mock) to add
+    items to the list of mock objects. Items added by these methods take precedence over the default items.
     """
     original_popen = subprocess.Popen
     original_mkdir = fileutil.mkdir
     original_path_exists = os.path.exists
     original_open = open
 
+    mocked_commands = _MOCKED_COMMANDS[:]
+    mocked_files = _MOCKED_FILES[:]
+    mocked_paths = _MOCKED_PATHS[:]
+
     def add_command_mock(pattern, output):
-        patcher.__commands.insert(0, (pattern, output))
+        mocked_commands.insert(0, (pattern, output))
+
+    def add_file_mock(actual, mock):
+        mocked_files.insert(0, (actual, mock))
+
+    def add_path_mock(mock):
+        mocked_paths.insert(0, mock)
 
     def mock_popen(command, *args, **kwargs):
         if isinstance(command, list):
@@ -119,7 +141,7 @@ def mock_cgroup_commands(tmp_dir):
         else:
             command_string = command
 
-        for cmd in patcher.__commands:
+        for cmd in mocked_commands:
             match = re.match(cmd[0], command_string)
             if match is not None:
                 command = ["echo", cmd[1]]
@@ -128,12 +150,12 @@ def mock_cgroup_commands(tmp_dir):
         return original_popen(command, *args, **kwargs)
 
     def get_mapped_path(path):
-        for item in __MOCKED_FILES:
+        for item in mocked_files:
             match = re.match(item[0], path)
             if match is not None:
                 return item[1]
 
-        for item in __MOCKED_PATHS:
+        for item in mocked_paths:
             mapped = re.sub(item, r"{0}\1".format(tmp_dir), path)
             if mapped != path:
                 mapped_parent = os.path.split(mapped)[0]
@@ -151,14 +173,34 @@ def mock_cgroup_commands(tmp_dir):
     def mock_path_exists(path):
         return original_path_exists(get_mapped_path(path))
 
+    data_files = [
+        (os.path.join(data_dir, 'init', 'walinuxagent.service'), UnitFilePaths.walinuxagent),
+        (os.path.join(data_dir, 'init', 'azure.slice'), UnitFilePaths.azure),
+        (os.path.join(data_dir, 'init', 'azure-vmextensions.slice'), UnitFilePaths.vmextensions)
+    ]
+
+    def add_data_file(source, target):
+        shutil.copyfile(source, get_mapped_path(target))
+
+    for items in data_files:
+        add_data_file(items[0], items[1])
+
     builtin_popen = "__builtin__.open" if sys.version_info[0] == 2 else "builtins.open"
-    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen) as patcher:
+
+    with patch("azurelinuxagent.common.cgroupapi.subprocess.Popen", side_effect=mock_popen):
         with patch("azurelinuxagent.common.cgroupconfigurator.fileutil.mkdir", side_effect=mock_mkdir):
             with patch("azurelinuxagent.common.cgroupapi.os.path.exists", side_effect=mock_path_exists):
                 with patch(builtin_popen, side_effect=mock_open):
                     with patch('azurelinuxagent.common.cgroupapi.CGroupsApi.cgroups_supported', return_value=True):
                         with patch('azurelinuxagent.common.cgroupapi.CGroupsApi.is_systemd', return_value=True):
-                            patcher.__commands = __MOCKED_COMMANDS[:]
-                            patcher.add_command_mock = add_command_mock
-                            patcher.get_mapped_path = get_mapped_path
-                            yield patcher
+                            class Mock:
+                                def __init__(self):
+                                    self.__commands = mocked_commands
+                                    self.__files = mocked_files
+                                    self.__paths = mocked_paths
+                                    self.add_command_mock = add_command_mock
+                                    self.add_file_mock = add_file_mock
+                                    self.add_path_mock = add_path_mock
+                                    self.add_data_file = add_data_file
+                                    self.get_mapped_path = get_mapped_path
+                            yield Mock()
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index cec4c2b0..1319d1cf 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -34,10 +34,9 @@ from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.event import WALAEventOperation
 from azurelinuxagent.common.exception import CGroupsException, ExtensionError, ExtensionErrorCodes
 from azurelinuxagent.common.future import ustr
-from azurelinuxagent.common.osutil import get_osutil
 from azurelinuxagent.common.utils import shellutil
-from tests.common.mock_cgroup_commands import mock_cgroup_commands
-from tests.tools import AgentTestCase, patch, mock_sleep, i_am_root
+from tests.common.mock_cgroup_commands import mock_cgroup_commands, UnitFilePaths
+from tests.tools import AgentTestCase, patch, mock_sleep, i_am_root, data_dir
 from tests.utils.miscellaneous_tools import format_processes, wait_for
 
 
@@ -71,8 +70,6 @@ class CGroupConfiguratorSystemdTestCase(AgentTestCase):
             self.assertTrue(configurator.enabled(), "Cgroups should be enabled")
             self.assertTrue(any(cg for cg in tracked if cg.name == 'walinuxagent.service' and 'cpu' in cg.path),
                 "The Agent's CPU is not being tracked. Tracked: {0}".format(tracked))
-            self.assertTrue(any(cg for cg in tracked if cg.name == 'walinuxagent.service' and 'memory' in cg.path),
-                "The Agent's memory is not being tracked. Tracked: {0}".format(tracked))
 
     def test_initialize_should_start_tracking_other_controllers_when_one_is_not_present(self):
         command_mocks = [(
@@ -84,9 +81,9 @@ cgroup on /sys/fs/cgroup/net_cls,net_prio type cgroup (rw,nosuid,nodev,noexec,re
 cgroup on /sys/fs/cgroup/perf_event type cgroup (rw,nosuid,nodev,noexec,relatime,perf_event)
 cgroup on /sys/fs/cgroup/hugetlb type cgroup (rw,nosuid,nodev,noexec,relatime,hugetlb)
 cgroup on /sys/fs/cgroup/freezer type cgroup (rw,nosuid,nodev,noexec,relatime,freezer)
-cgroup on /sys/fs/cgroup/memory type cgroup (rw,nosuid,nodev,noexec,relatime,memory)
 cgroup on /sys/fs/cgroup/pids type cgroup (rw,nosuid,nodev,noexec,relatime,pids)
 cgroup on /sys/fs/cgroup/devices type cgroup (rw,nosuid,nodev,noexec,relatime,devices)
+cgroup on /sys/fs/cgroup/cpu,cpuacct type cgroup (rw,nosuid,nodev,noexec,relatime,cpu,cpuacct)
 cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blkio)
 ''')]
         with self._get_cgroup_configurator(command_mocks=command_mocks) as configurator:
@@ -94,10 +91,8 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             tracked = CGroupsTelemetry._tracked  # pylint: disable=protected-access
 
             self.assertTrue(configurator.enabled(), "Cgroups should be enabled")
-            self.assertFalse(any(cg for cg in tracked if cg.name == 'walinuxagent.service' and 'cpu' in cg.path),
-                "The Agent's CPU should not be tracked. Tracked: {0}".format(tracked))
-            self.assertTrue(any(cg for cg in tracked if cg.name == 'walinuxagent.service' and 'memory' in cg.path),
-                "The Agent's memory is not being tracked. Tracked: {0}".format(tracked))
+            self.assertFalse(any(cg for cg in tracked if cg.name == 'walinuxagent.service' and 'memory' in cg.path),
+                "The Agent's memory should not be tracked. Tracked: {0}".format(tracked))
 
     def test_initialize_should_not_enable_cgroups_is_the_cpu_and_memory_controllers_are_not_present(self):
         command_mocks = [(
@@ -142,20 +137,50 @@ cgroup on /sys/fs/cgroup/blkio type cgroup (rw,nosuid,nodev,noexec,relatime,blki
             self.assertFalse(configurator.enabled(), "Cgroups should not be enabled")
             self.assertEqual(len(tracked), 0, "No cgroups should be tracked. Tracked: {0}".format(tracked))
 
-    def test_initialize_should_create_slice_unit_files_when_they_do_not_exist(self):
+    def test_initialize_should_not_create_unit_files(self):
         with self._get_cgroup_configurator() as configurator:
-            # When the unit files do not exist under /lib/systemd/system, initialize() should create them
-            # (note that the mock CGroupConfigurator maps these files to a test location, so we need to use the mapped path)
-            agent_unit_name = get_osutil().get_service_name() + ".service"
-            azure_slice_unit_file = configurator.mocks.get_mapped_path("/lib/systemd/system/azure.slice")
-            extensions_slice_unit_file = configurator.mocks.get_mapped_path("/lib/systemd/system/azure-vmextensions.slice")
-            agent_drop_in_file_slice = configurator.mocks.get_mapped_path("/lib/systemd/system/{0}.d/10-Slice.conf".format(agent_unit_name))
-            agent_drop_in_file_cpu_accounting = configurator.mocks.get_mapped_path("/lib/systemd/system/{0}.d/11-CPUAccounting.conf".format(agent_unit_name))
-
+            # get the paths to the mocked files
+            azure_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.azure)
+            extensions_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.vmextensions)
+            agent_drop_in_file_slice = configurator.mocks.get_mapped_path(UnitFilePaths.slice)
+            agent_drop_in_file_cpu_accounting = configurator.mocks.get_mapped_path(UnitFilePaths.cpu_accounting)
+            agent_drop_in_file_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.cpu_quota)
+
+            # The mock creates the slice unit files; delete them
+            os.remove(azure_slice_unit_file)
+            os.remove(extensions_slice_unit_file)
+
+            # The service file for the agent includes settings for the slice and cpu accounting, but not for cpu quota; initialize()
+            # should not create drop in files for the first 2, but it should create one the cpu quota
+            self.assertFalse(os.path.exists(azure_slice_unit_file), "{0} should not have been created".format(azure_slice_unit_file))
+            self.assertFalse(os.path.exists(extensions_slice_unit_file), "{0} should not have been created".format(extensions_slice_unit_file))
+            self.assertFalse(os.path.exists(agent_drop_in_file_slice), "{0} should not have been created".format(agent_drop_in_file_slice))
+            self.assertFalse(os.path.exists(agent_drop_in_file_cpu_accounting), "{0} should not have been created".format(agent_drop_in_file_cpu_accounting))
+            self.assertTrue(os.path.exists(agent_drop_in_file_cpu_quota), "{0} was not created".format(agent_drop_in_file_cpu_quota))
+
+    def test_initialize_should_create_unit_files_when_the_agent_service_file_is_not_updated(self):
+        with self._get_cgroup_configurator(initialize=False) as configurator:
+            # get the paths to the mocked files
+            azure_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.azure)
+            extensions_slice_unit_file = configurator.mocks.get_mapped_path(UnitFilePaths.vmextensions)
+            agent_drop_in_file_slice = configurator.mocks.get_mapped_path(UnitFilePaths.slice)
+            agent_drop_in_file_cpu_accounting = configurator.mocks.get_mapped_path(UnitFilePaths.cpu_accounting)
+            agent_drop_in_file_cpu_quota = configurator.mocks.get_mapped_path(UnitFilePaths.cpu_quota)
+
+            # The mock creates the service and slice unit files; replace the former and delete the latter
+            configurator.mocks.add_data_file(os.path.join(data_dir, 'init', "walinuxagent.service.previous"), UnitFilePaths.walinuxagent)
+            os.remove(azure_slice_unit_file)
+            os.remove(extensions_slice_unit_file)
+
+            configurator.initialize()
+
+            # The older service file for the agent did not include settings for the slice and cpu parameters; in that case, initialize() should
+            # create drop in files to set those properties
             self.assertTrue(os.path.exists(azure_slice_unit_file), "{0} was not created".format(azure_slice_unit_file))
             self.assertTrue(os.path.exists(extensions_slice_unit_file), "{0} was not created".format(extensions_slice_unit_file))
             self.assertTrue(os.path.exists(agent_drop_in_file_slice), "{0} was not created".format(agent_drop_in_file_slice))
             self.assertTrue(os.path.exists(agent_drop_in_file_cpu_accounting), "{0} was not created".format(agent_drop_in_file_cpu_accounting))
+            self.assertTrue(os.path.exists(agent_drop_in_file_cpu_quota), "{0} was not created".format(agent_drop_in_file_cpu_quota))
 
     def test_enable_and_disable_should_change_the_enabled_state_of_cgroups(self):
         with self._get_cgroup_configurator() as configurator:
diff --git a/tests/data/init/azure-vmextensions.slice b/tests/data/init/azure-vmextensions.slice
new file mode 100644
index 00000000..87dd84b1
--- /dev/null
+++ b/tests/data/init/azure-vmextensions.slice
@@ -0,0 +1,6 @@
+[Unit]
+Description=Slice for Azure VM Extensions
+DefaultDependencies=no
+Before=slices.target
+[Slice]
+CPUAccounting=yes
diff --git a/tests/data/init/azure.slice b/tests/data/init/azure.slice
new file mode 100644
index 00000000..6310d192
--- /dev/null
+++ b/tests/data/init/azure.slice
@@ -0,0 +1,4 @@
+[Unit]
+Description=Slice for Azure VM Agent and Extensions
+DefaultDependencies=no
+Before=slices.target
diff --git a/tests/data/init/walinuxagent.service b/tests/data/init/walinuxagent.service
new file mode 100644
index 00000000..cc1b71c9
--- /dev/null
+++ b/tests/data/init/walinuxagent.service
@@ -0,0 +1,22 @@
+#
+# NOTE: This is the service file used in current versions of the agent (>= 2.2.55)
+#
+[Unit]
+Description=Azure Linux Agent
+
+After=network-online.target cloud-init.service
+Wants=network-online.target sshd.service sshd-keygen.service
+
+ConditionFileIsExecutable=/usr/sbin/waagent
+ConditionPathExists=/etc/waagent.conf
+
+[Service]
+Type=simple
+ExecStart=/usr/bin/python3 -u /usr/sbin/waagent -daemon
+Restart=always
+Slice=azure.slice
+CPUAccounting=yes
+
+[Install]
+WantedBy=multi-user.target
+
diff --git a/tests/data/init/walinuxagent.service.previous b/tests/data/init/walinuxagent.service.previous
new file mode 100644
index 00000000..297c6a98
--- /dev/null
+++ b/tests/data/init/walinuxagent.service.previous
@@ -0,0 +1,20 @@
+#
+# NOTE: This is the service file used in older versions of the agent (<= 2.2.54)
+#
+[Unit]
+Description=Azure Linux Agent
+
+After=network-online.target cloud-init.service
+Wants=network-online.target sshd.service sshd-keygen.service
+
+ConditionFileIsExecutable=/usr/sbin/waagent
+ConditionPathExists=/etc/waagent.conf
+
+[Service]
+Type=simple
+ExecStart=/usr/bin/python3 -u /usr/sbin/waagent -daemon
+Restart=always
+
+[Install]
+WantedBy=multi-user.target
+