Repository: Azure/WALinuxAgent
Pull request #2671 (2022-09-27 23:28:24)

enforce memory usage for agent
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description
Instead enforcing memory with cgroup, agent itself monitor the current memory usage at regular intervals and do the safe exit when memory reaches the limit. This approach avoids OOM kills on system.

By default, memory usage check set to false for now. We will enable once we know the agent memory limit value.

Issue # <!-- if any -->
<!--
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).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

codecov[bot]: # [Codecov](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2671](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (bef1eac) into [develop](https://codecov.io/gh/Azure/WALinuxAgent/commit/45e960da74f6c1eda090841720c42bb1bf08f860?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (45e960d) will **increase** coverage by `0.03%`.
> The diff coverage is `87.50%`.

```diff
@@             Coverage Diff             @@
##           develop    #2671      +/-   ##
===========================================
+ Coverage    71.96%   71.99%   +0.03%     
===========================================
  Files          103      103              
  Lines        15679    15717      +38     
  Branches      2486     2495       +9     
===========================================
+ Hits         11283    11316      +33     
  Misses        3879     3879              
- Partials       517      522       +5     
```


| [Impacted Files](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/cgroupconfigurator.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jZ3JvdXBjb25maWd1cmF0b3IucHk=) | `73.65% <78.57%> (+0.06%)` | :arrow_up: |
| [azurelinuxagent/ga/update.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL3VwZGF0ZS5weQ==) | `89.36% <90.90%> (+0.02%)` | :arrow_up: |
| [azurelinuxagent/common/conf.py](https://codecov.io/gh/Azure/WALinuxAgent/pull/2671/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9jb25mLnB5) | `79.83% <100.00%> (+0.34%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)


Patch:
diff --git a/azurelinuxagent/common/cgroupconfigurator.py b/azurelinuxagent/common/cgroupconfigurator.py
index 47f1da35..0abee220 100644
--- a/azurelinuxagent/common/cgroupconfigurator.py
+++ b/azurelinuxagent/common/cgroupconfigurator.py
@@ -26,7 +26,7 @@ from azurelinuxagent.common import logger
 from azurelinuxagent.common.cgroup import CpuCgroup, AGENT_NAME_TELEMETRY, MetricsCounter, MemoryCgroup
 from azurelinuxagent.common.cgroupapi import CGroupsApi, SystemdCgroupsApi, SystemdRunError, EXTENSION_SLICE_PREFIX
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
-from azurelinuxagent.common.exception import ExtensionErrorCodes, CGroupsException
+from azurelinuxagent.common.exception import ExtensionErrorCodes, CGroupsException, AgentMemoryExceededException
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil, systemd
 from azurelinuxagent.common.version import get_distro
@@ -143,6 +143,7 @@ class CGroupConfigurator(object):
             self._cgroups_api = None
             self._agent_cpu_cgroup_path = None
             self._agent_memory_cgroup_path = None
+            self._agent_memory_cgroup = None
             self._check_cgroups_lock = threading.RLock()  # Protect the check_cgroups which is called from Monitor thread and main loop.
 
         def initialize(self):
@@ -194,7 +195,8 @@ class CGroupConfigurator(object):
 
                 if self._agent_memory_cgroup_path is not None:
                     _log_cgroup_info("Agent Memory cgroup: {0}", self._agent_memory_cgroup_path)
-                    CGroupsTelemetry.track_cgroup(MemoryCgroup(AGENT_NAME_TELEMETRY, self._agent_memory_cgroup_path))
+                    self._agent_memory_cgroup = MemoryCgroup(AGENT_NAME_TELEMETRY, self._agent_memory_cgroup_path)
+                    CGroupsTelemetry.track_cgroup(self._agent_memory_cgroup)
 
                 _log_cgroup_info('Agent cgroups enabled: {0}', self._agent_cgroups_enabled)
 
@@ -729,6 +731,19 @@ class CGroupConfigurator(object):
                     if metric.value > conf.get_agent_cpu_throttled_time_threshold():
                         raise CGroupsException("The agent has been throttled for {0} seconds".format(metric.value))
 
+        def check_agent_memory_usage(self):
+            if self.enabled() and self._agent_memory_cgroup:
+                metrics = self._agent_memory_cgroup.get_tracked_metrics()
+                current_usage = 0
+                for metric in metrics:
+                    if metric.counter == MetricsCounter.TOTAL_MEM_USAGE:
+                        current_usage += metric.value
+                    elif metric.counter == MetricsCounter.SWAP_MEM_USAGE:
+                        current_usage += metric.value
+
+                if current_usage > conf.get_agent_memory_quota():
+                    raise AgentMemoryExceededException("The agent memory limit {0} bytes exceeded. The current reported usage is {1} bytes.".format(conf.get_agent_memory_quota(), current_usage))
+
         @staticmethod
         def _get_parent(pid):
             """
diff --git a/azurelinuxagent/common/conf.py b/azurelinuxagent/common/conf.py
index 3c6e960f..46765ea9 100644
--- a/azurelinuxagent/common/conf.py
+++ b/azurelinuxagent/common/conf.py
@@ -136,6 +136,7 @@ __SWITCH_OPTIONS__ = {
     "Debug.CgroupLogMetrics": False,
     "Debug.CgroupDisableOnProcessCheckFailure": True,
     "Debug.CgroupDisableOnQuotaCheckFailure": True,
+    "Debug.EnableAgentMemoryUsageCheck": False,
     "Debug.EnableFastTrack": True,
     "Debug.EnableGAVersioning": False
 }
@@ -186,6 +187,7 @@ __INTEGER_OPTIONS__ = {
     "Debug.CgroupCheckPeriod": 300,
     "Debug.AgentCpuQuota": 50,
     "Debug.AgentCpuThrottledTimeThreshold": 120,
+    "Debug.AgentMemoryQuota": 30 * 1024 ** 2,
     "Debug.EtpCollectionPeriod": 300,
     "Debug.AutoUpdateHotfixFrequency": 14400,
     "Debug.AutoUpdateNormalFrequency": 86400,
@@ -555,6 +557,24 @@ def get_agent_cpu_throttled_time_threshold(conf=__conf__):
     return conf.get_int("Debug.AgentCpuThrottledTimeThreshold", 120)
 
 
+def get_agent_memory_quota(conf=__conf__):
+    """
+    Memory quota for the agent in bytes.
+
+    NOTE: This option is experimental and may be removed in later versions of the Agent.
+    """
+    return conf.get_int("Debug.AgentMemoryQuota", 30 * 1024 ** 2)
+
+
+def get_enable_agent_memory_usage_check(conf=__conf__):
+    """
+    If True, Agent checks it's Memory usage.
+
+    NOTE: This option is experimental and may be removed in later versions of the Agent.
+    """
+    return conf.get_switch("Debug.EnableAgentMemoryUsageCheck", False)
+
+
 def get_cgroup_monitor_expiry_time(conf=__conf__):
     """
     cgroups monitoring for pilot extensions disabled after expiry time
diff --git a/azurelinuxagent/common/event.py b/azurelinuxagent/common/event.py
index b7aba5e4..1f903a9f 100644
--- a/azurelinuxagent/common/event.py
+++ b/azurelinuxagent/common/event.py
@@ -69,6 +69,7 @@ class WALAEventOperation:
     ActivateResourceDisk = "ActivateResourceDisk"
     AgentBlacklisted = "AgentBlacklisted"
     AgentEnabled = "AgentEnabled"
+    AgentMemory = "AgentMemory"
     AgentUpgrade = "AgentUpgrade"
     ArtifactsProfileBlob = "ArtifactsProfileBlob"
     CGroupsCleanUp = "CGroupsCleanUp"
diff --git a/azurelinuxagent/common/exception.py b/azurelinuxagent/common/exception.py
index 9b16c426..04846623 100644
--- a/azurelinuxagent/common/exception.py
+++ b/azurelinuxagent/common/exception.py
@@ -58,6 +58,14 @@ class AgentConfigError(AgentError):
         super(AgentConfigError, self).__init__(msg, inner)
 
 
+class AgentMemoryExceededException(AgentError):
+    """
+    When Agent memory limit reached.
+    """
+    def __init__(self, msg=None, inner=None):
+        super(AgentMemoryExceededException, self).__init__(msg, inner)
+
+
 class AgentNetworkError(AgentError):
     """
     When network is not available.
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 0fac66f2..58766847 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -39,7 +39,8 @@ from azurelinuxagent.common.agent_supported_feature import get_supported_feature
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator
 from azurelinuxagent.common.event import add_event, initialize_event_logger_vminfo_common_parameters, \
     WALAEventOperation, EVENTS_DIRECTORY
-from azurelinuxagent.common.exception import ResourceGoneError, UpdateError, ExitException, AgentUpgradeExitException
+from azurelinuxagent.common.exception import ResourceGoneError, UpdateError, ExitException, AgentUpgradeExitException, \
+    AgentMemoryExceededException
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.osutil import get_osutil, systemd
 from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
@@ -137,6 +138,7 @@ def get_update_handler():
 
 class UpdateHandler(object):
     TELEMETRY_HEARTBEAT_PERIOD = timedelta(minutes=30)
+    CHECK_MEMORY_USAGE_PERIOD = timedelta(seconds=conf.get_cgroup_check_period())
 
     def __init__(self):
         self.osutil = get_osutil()
@@ -162,6 +164,9 @@ class UpdateHandler(object):
         self._heartbeat_id = str(uuid.uuid4()).upper()
         self._heartbeat_counter = 0
 
+        self._last_check_memory_usage = datetime.min
+        self._check_memory_usage_last_error_report = datetime.min
+
         # VM Size is reported via the heartbeat, default it here.
         self._vm_size = None
 
@@ -401,6 +406,7 @@ class UpdateHandler(object):
                 self._check_threads_running(all_thread_handlers)
                 self._process_goal_state(exthandlers_handler, remote_access_handler)
                 self._send_heartbeat_telemetry(protocol)
+                self._check_agent_memory_usage()
                 time.sleep(self._goal_state_period)
 
         except AgentUpgradeExitException as exitException:
@@ -1288,6 +1294,27 @@ class UpdateHandler(object):
             self._heartbeat_update_goal_state_error_count = 0
             self._last_telemetry_heartbeat = datetime.utcnow()
 
+    def _check_agent_memory_usage(self):
+        """
+        This checks the agent current memory usage and safely exit the process if agent reaches the memory limit
+        """
+        try:
+            if conf.get_enable_agent_memory_usage_check() and self._extensions_summary.converged:
+                if self._last_check_memory_usage == datetime.min or datetime.utcnow() >= (self._last_check_memory_usage + UpdateHandler.CHECK_MEMORY_USAGE_PERIOD):
+                    self._last_check_memory_usage = datetime.utcnow()
+                    CGroupConfigurator.get_instance().check_agent_memory_usage()
+        except AgentMemoryExceededException as exception:
+            msg = "Check on agent memory usage:\n{0}".format(ustr(exception))
+            logger.info(msg)
+            add_event(AGENT_NAME, op=WALAEventOperation.AgentMemory, is_success=True, message=msg)
+            raise ExitException("Agent {0} is reached memory limit -- exiting".format(CURRENT_AGENT))
+        except Exception as exception:
+            if self._check_memory_usage_last_error_report == datetime.min or (self._check_memory_usage_last_error_report + timedelta(hours=6)) > datetime.now():
+                self._check_memory_usage_last_error_report = datetime.now()
+                msg = "Error checking the agent's memory usage: {0} --- [NOTE: Will not log the same error for the 6 hours]".format(ustr(exception))
+                logger.warn(msg)
+                add_event(AGENT_NAME, op=WALAEventOperation.AgentMemory, is_success=False, message=msg)
+
     @staticmethod
     def _ensure_extension_telemetry_state_configured_properly(protocol):
         etp_enabled = get_supported_feature_by_name(SupportedFeatureNames.ExtensionTelemetryPipeline).is_supported


Test patch:
diff --git a/tests/common/test_cgroupconfigurator.py b/tests/common/test_cgroupconfigurator.py
index 62a7211a..60a7cfde 100644
--- a/tests/common/test_cgroupconfigurator.py
+++ b/tests/common/test_cgroupconfigurator.py
@@ -33,7 +33,8 @@ from azurelinuxagent.common.cgroup import AGENT_NAME_TELEMETRY, MetricsCounter,
 from azurelinuxagent.common.cgroupconfigurator import CGroupConfigurator, DisableCgroups
 from azurelinuxagent.common.cgroupstelemetry import CGroupsTelemetry
 from azurelinuxagent.common.event import WALAEventOperation
-from azurelinuxagent.common.exception import CGroupsException, ExtensionError, ExtensionErrorCodes
+from azurelinuxagent.common.exception import CGroupsException, ExtensionError, ExtensionErrorCodes, \
+    AgentMemoryExceededException
 from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.utils import shellutil, fileutil
 from tests.common.mock_environment import MockCommand
@@ -986,3 +987,15 @@ exit 0
                 finally:
                     for p in patchers:
                         p.stop()
+
+    def test_check_agent_memory_usage_should_raise_a_cgroups_exception_when_the_limit_is_exceeded(self):
+        metrics = [MetricValue(MetricsCategory.MEMORY_CATEGORY, MetricsCounter.TOTAL_MEM_USAGE, AGENT_NAME_TELEMETRY, conf.get_agent_memory_quota() + 1),
+                   MetricValue(MetricsCategory.MEMORY_CATEGORY, MetricsCounter.SWAP_MEM_USAGE, AGENT_NAME_TELEMETRY, conf.get_agent_memory_quota() + 1)]
+
+        with self.assertRaises(AgentMemoryExceededException) as context_manager:
+            with self._get_cgroup_configurator() as configurator:
+                with patch("azurelinuxagent.common.cgroup.MemoryCgroup.get_tracked_metrics") as tracked_metrics:
+                    tracked_metrics.return_value = metrics
+                    configurator.check_agent_memory_usage()
+
+        self.assertIn("The agent memory limit {0} bytes exceeded".format(conf.get_agent_memory_quota()), ustr(context_manager.exception), "An incorrect exception was raised")
\ No newline at end of file
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 695b6d57..0bbd2279 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -28,7 +28,8 @@ _ORIGINAL_POPEN = subprocess.Popen
 
 from azurelinuxagent.common import conf
 from azurelinuxagent.common.event import EVENTS_DIRECTORY, WALAEventOperation
-from azurelinuxagent.common.exception import ProtocolError, UpdateError, ResourceGoneError, HttpError
+from azurelinuxagent.common.exception import ProtocolError, UpdateError, ResourceGoneError, HttpError, \
+    ExitException, AgentMemoryExceededException
 from azurelinuxagent.common.future import ustr, httpclient
 from azurelinuxagent.common.persist_firewall_rules import PersistFirewallRulesHandler
 from azurelinuxagent.common.protocol.hostplugin import URI_FORMAT_GET_API_VERSIONS, HOST_PLUGIN_PORT, \
@@ -1461,7 +1462,7 @@ class TestUpdate(UpdateTestCase):
         eh = Extension(name=name)
         eh.version = version
         return ExtHandlerInstance(eh, protocol)
-    
+
     def test_update_handler_recovers_from_error_with_no_certs(self):
         data = DATA_FILE.copy()
         data['goal_state'] = 'wire/goal_state_no_certs.xml'
@@ -1489,7 +1490,7 @@ class TestUpdate(UpdateTestCase):
                             for (args, _) in filter(lambda a: len(a) > 0, patched_error.call_args_list):
                                 if unexpected_msg_fragment in args[0]:
                                     matching_errors.append(args[0])
-                            
+
                             if len(matching_errors) > 1:
                                 self.fail("Guest Agent did not recover, with new error(s): {}"\
                                     .format(matching_errors[1:]))
@@ -2894,7 +2895,7 @@ class ProcessGoalStateTestCase(AgentTestCase):
             if HttpRequestPredicates.is_host_plugin_vm_settings_request(url):
                 return MockHttpResponse(404)
             return None
-        
+
         with mock_wire_protocol(data) as protocol:
 
             def mock_live_migration(iteration):
@@ -2904,7 +2905,7 @@ class ProcessGoalStateTestCase(AgentTestCase):
                 elif iteration == 2:
                     protocol.mock_wire_data.set_incarnation(3)
                     protocol.set_http_handlers(http_get_handler=vm_settings_not_supported)
-            
+
             with mock_update_handler(protocol, 3, on_new_iteration=mock_live_migration) as update_handler:
                 with patch("azurelinuxagent.ga.update.logger.error") as patched_error:
                     def check_for_errors():
@@ -2916,7 +2917,7 @@ class ProcessGoalStateTestCase(AgentTestCase):
 
                     update_handler.run(debug=True)
                     check_for_errors()
-                
+
             timestamp = protocol.client.get_host_plugin()._fast_track_timestamp
             self.assertEqual(timestamp, timeutil.create_timestamp(datetime.min),
                 "Expected fast track time stamp to be set to {0}, got {1}".format(datetime.min, timestamp))
@@ -2926,16 +2927,16 @@ class HeartbeatTestCase(AgentTestCase):
     @patch("azurelinuxagent.common.logger.info")
     @patch("azurelinuxagent.ga.update.add_event")
     def test_telemetry_heartbeat_creates_event(self, patch_add_event, patch_info, *_):
-        
+
         with mock_wire_protocol(mockwiredata.DATA_FILE) as mock_protocol:
             update_handler = get_update_handler()
-            
+
             update_handler.last_telemetry_heartbeat = datetime.utcnow() - timedelta(hours=1)
             update_handler._send_heartbeat_telemetry(mock_protocol)
             self.assertEqual(1, patch_add_event.call_count)
             self.assertTrue(any(call_args[0] == "[HEARTBEAT] Agent {0} is running as the goal state agent {1}"
                             for call_args in patch_info.call_args), "The heartbeat was not written to the agent's log")
-    
+
     @patch("azurelinuxagent.ga.update.add_event")
     @patch("azurelinuxagent.common.protocol.imds.ImdsClient")
     def test_telemetry_heartbeat_retries_failed_vm_size_fetch(self, mock_imds_factory, patch_add_event, *_):
@@ -2953,7 +2954,7 @@ class HeartbeatTestCase(AgentTestCase):
             self.assertTrue(telemetry_message.endswith(vm_size),
                 "Expected HeartBeat message ('{0}') to end with the test vmSize value, {1}."\
                 .format(telemetry_message, vm_size))
-        
+
         with mock_wire_protocol(mockwiredata.DATA_FILE) as mock_protocol:
             update_handler = get_update_handler()
             update_handler.protocol_util.get_protocol = Mock(return_value=mock_protocol)
@@ -2979,6 +2980,47 @@ class HeartbeatTestCase(AgentTestCase):
             validate_single_heartbeat_event_matches_vm_size("TestVmSizeValue")
 
 
+class AgentMemoryCheckTestCase(AgentTestCase):
+
+    @patch("azurelinuxagent.common.logger.info")
+    @patch("azurelinuxagent.ga.update.add_event")
+    def test_check_agent_memory_usage_raises_exit_exception(self, patch_add_event, patch_info, *_):
+        with patch("azurelinuxagent.common.cgroupconfigurator.CGroupConfigurator._Impl.check_agent_memory_usage", side_effect=AgentMemoryExceededException()):
+            with patch('azurelinuxagent.common.conf.get_enable_agent_memory_usage_check', return_value=True):
+                with self.assertRaises(ExitException) as context_manager:
+                    update_handler = get_update_handler()
+
+                    update_handler._check_agent_memory_usage()
+                    self.assertEqual(1, patch_add_event.call_count)
+                    self.assertTrue(any("Check on agent memory usage" in call_args[0]
+                                        for call_args in patch_info.call_args),
+                                    "The memory check was not written to the agent's log")
+                    self.assertIn("Agent {0} is reached memory limit -- exiting".format(CURRENT_AGENT),
+                                  ustr(context_manager.exception), "An incorrect exception was raised")
+
+    @patch("azurelinuxagent.common.logger.warn")
+    @patch("azurelinuxagent.ga.update.add_event")
+    def test_check_agent_memory_usage_fails(self, patch_add_event, patch_warn, *_):
+        with patch("azurelinuxagent.common.cgroupconfigurator.CGroupConfigurator._Impl.check_agent_memory_usage", side_effect=Exception()):
+            with patch('azurelinuxagent.common.conf.get_enable_agent_memory_usage_check', return_value=True):
+                update_handler = get_update_handler()
+
+                update_handler._check_agent_memory_usage()
+                self.assertTrue(any("Error checking the agent's memory usage" in call_args[0]
+                                    for call_args in patch_warn.call_args),
+                                "The memory check was not written to the agent's log")
+                self.assertEqual(1, patch_add_event.call_count)
+                add_events = [kwargs for _, kwargs in patch_add_event.call_args_list if
+                                kwargs["op"] == WALAEventOperation.AgentMemory]
+                self.assertTrue(
+                    len(add_events) == 1,
+                    "Exactly 1 event should have been emitted when memory usage check fails. Got: {0}".format(add_events))
+                self.assertIn(
+                    "Error checking the agent's memory usage",
+                    add_events[0]["message"],
+                    "The error message is not correct when memory usage check failed")
+
+
 class GoalStateIntervalTestCase(AgentTestCase):
     def test_initial_goal_state_period_should_default_to_goal_state_period(self):
         configuration_provider = conf.ConfigurationProvider()
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 1b14c9d1..c585e845 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -31,6 +31,7 @@ Autoupdate.Frequency = 3600
 DVD.MountPoint = /mnt/cdrom/secure
 Debug.AgentCpuQuota = 50
 Debug.AgentCpuThrottledTimeThreshold = 120
+Debug.AgentMemoryQuota = 31457280
 Debug.AutoUpdateHotfixFrequency = 14400
 Debug.AutoUpdateNormalFrequency = 86400
 Debug.CgroupCheckPeriod = 300
@@ -39,6 +40,7 @@ Debug.CgroupDisableOnQuotaCheckFailure = True
 Debug.CgroupLogMetrics = False
 Debug.CgroupMonitorExpiryTime = 2022-03-31
 Debug.CgroupMonitorExtensionName = Microsoft.Azure.Monitor.AzureMonitorLinuxAgent
+Debug.EnableAgentMemoryUsageCheck = False
 Debug.EnableFastTrack = True
 Debug.EnableGAVersioning = False
 Debug.EtpCollectionPeriod = 300