Skip to content

Commit

Permalink
Address test comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@mgunnala
mgunnala committed Dec 6, 2024
1 parent b037e41 commit ba3869c
Show file tree
Hide file tree
Showing 6 changed files with 236 additions and 231 deletions.
64 changes: 36 additions & 28 deletions azurelinuxagent/ga/exthandlers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -528,14 +528,16 @@ def handle_ext_handlers(self, goal_state_id):
logger.info("{0}: {1}".format(ext_full_name, msg))
add_event(op=WALAEventOperation.ExtensionProcessing, message="{0}: {1}".format(ext_full_name, msg))
handler_i.set_handler_status(status=ExtHandlerStatusValue.not_ready, message=msg, code=-1)
handler_i.create_status_file_if_not_exist(extension,
status=ExtensionStatusValue.error,
code=-1,
operation=handler_i.operation,
message=msg)
handler_i.create_status_file(extension,
status=ExtensionStatusValue.error,
code=-1,
operation=handler_i.operation,
message=msg)
continue
# Invoke policy engine to determine if extension is allowed. If not, block extension and report error on
# behalf of the extension.

# If an error was thrown during policy engine initialization, skip further processing of the extension.
# CRP is still waiting for status, so we report error status here.
# of the extension.
policy_op, policy_err_code = _POLICY_ERROR_MAP.get(ext_handler.state)
if policy_error is not None:
msg = "Extension will not be processed: {0}".format(ustr(policy_error))
Expand All @@ -544,15 +546,6 @@ def handle_ext_handlers(self, goal_state_id):
extension=extension)
continue

extension_allowed = policy_engine.should_allow_extension(ext_handler.name)
if not extension_allowed:
msg = (
"Extension will not be processed: failed to {0} extension '{1}' because it is not specified "
"in the allowlist. To {0}, add the extension to the allowed list in the policy file ('{2}')."
).format(policy_op, ext_handler.name, conf.get_policy_file_path())
self.__report_policy_error(handler_i, policy_err_code, report_op=handler_i.operation,
message=msg, extension=extension)

# In case of depends-on errors, we skip processing extensions if there was an error processing dependent extensions.
# But CRP is still waiting for some status back for the skipped extensions. In order to propagate the status back to CRP,
# we will report status back here with the relevant error message for each of the dependent extension.
Expand All @@ -564,16 +557,27 @@ def handle_ext_handlers(self, goal_state_id):
if handler_i.get_handler_status() is None:
handler_i.set_handler_status(message=depends_on_err_msg, code=-1)

handler_i.create_status_file_if_not_exist(extension, status=ExtensionStatusValue.error, code=-1,
operation=WALAEventOperation.ExtensionProcessing,
message=depends_on_err_msg)
handler_i.create_status_file(extension, status=ExtensionStatusValue.error, code=-1,
operation=WALAEventOperation.ExtensionProcessing,
message=depends_on_err_msg)

# For SC extensions, overwrite the HandlerStatus with the relevant message
else:
handler_i.set_handler_status(message=depends_on_err_msg, code=-1)

continue

# Invoke policy engine to determine if extension is allowed. If disallowed, report an error on behalf of
# the extension and do not process the extension. Dependent extensions will also be blocked.
extension_allowed = policy_engine.should_allow_extension(ext_handler.name)
if not extension_allowed:
msg = (
"Extension will not be processed: failed to {0} extension '{1}' because it is not specified "
"in the allowlist. To {0}, add the extension to the allowed list in the policy file ('{2}')."
).format(policy_op, ext_handler.name, conf.get_policy_file_path())
self.__report_policy_error(handler_i, policy_err_code, report_op=handler_i.operation,
message=msg, extension=extension)

# Process extensions and get if it was successfully executed or not
# If extension was blocked by policy, treat the extension as failed and do not process the handler.
if not extension_allowed:
Expand Down Expand Up @@ -694,8 +698,8 @@ def handle_ext_handler(self, ext_handler_i, extension, goal_state_id):
# This error is only thrown for enable operation on MultiConfig extension.
# Since these are maintained by the extensions, the expectation here is that they would update their status files appropriately with their errors.
# The extensions should already have a placeholder status file, but incase they dont, setting one here to fail fast.
ext_handler_i.create_status_file_if_not_exist(extension, status=ExtensionStatusValue.error, code=error.code,
operation=ext_handler_i.operation, message=err_msg)
ext_handler_i.create_status_file(extension, status=ExtensionStatusValue.error, code=error.code,
operation=ext_handler_i.operation, message=err_msg)
add_event(name=ext_name, version=ext_handler_i.ext_handler.version, op=ext_handler_i.operation,
is_success=False, log_event=True, message=err_msg)
except ExtensionsGoalStateError as error:
Expand Down Expand Up @@ -735,8 +739,8 @@ def __handle_and_report_ext_handler_errors(ext_handler_i, error, report_op, mess
# file with failure since the extensions wont be called where they can create their status files.
# This way we guarantee reporting back to CRP
if ext_handler_i.should_perform_multi_config_op(extension):
ext_handler_i.create_status_file_if_not_exist(extension, status=ExtensionStatusValue.error, code=error.code,
operation=report_op, message=message)
ext_handler_i.create_status_file(extension, status=ExtensionStatusValue.error, code=error.code,
operation=report_op, message=message)

if report:
name = ext_handler_i.get_extension_full_name(extension)
Expand All @@ -760,9 +764,11 @@ def __report_policy_error(ext_handler_i, error_code, report_op, message, extensi
ext_handler_i.set_handler_status(message=message, code=error_code)

# Create status file for extensions with settings (single and multi config).
# If status file already exists, overwrite it. If an extension was previously reporting status and is now
# blocked by a policy error, we should report the policy error.
if extension is not None:
ext_handler_i.create_status_file_if_not_exist(extension, status=ExtensionStatusValue.error, code=error_code,
operation=report_op, message=message)
ext_handler_i.create_status_file(extension, status=ExtensionStatusValue.error, code=error_code,
operation=report_op, message=message, overwrite=True)

name = ext_handler_i.get_extension_full_name(extension)
handler_version = ext_handler_i.ext_handler.version
Expand Down Expand Up @@ -1067,7 +1073,7 @@ def report_ext_handler_status(self, vm_status, ext_handler, goal_state_changed):
# extension even if HandlerState == NotInstalled (Sample scenario: ExtensionsGoalStateError, DecideVersionError, etc)
# We also need to report extension status for an uninstalled handler if extensions are disabled, or if the extension
# failed due to policy, because CRP waits for extension runtime status before failing the extension operation.
if handler_state != ExtHandlerState.NotInstalled or ext_handler.supports_multi_config or not conf.get_extensions_enabled() or ExtensionPolicyEngine.get_policy_enforcement_enabled():
if handler_state != ExtHandlerState.NotInstalled or ext_handler.supports_multi_config or not conf.get_extensions_enabled():

# Since we require reading the Manifest for reading the heartbeat, this would fail if HandlerManifest not found.
# Only try to read heartbeat if HandlerState != NotInstalled.
Expand Down Expand Up @@ -1420,9 +1426,11 @@ def set_extension_resource_limits(self):
extension_name=extension_name, cpu_quota=resource_limits.get_extension_slice_cpu_quota())
CGroupConfigurator.get_instance().set_extension_services_cpu_memory_quota(resource_limits.get_service_list())

def create_status_file_if_not_exist(self, extension, status, code, operation, message):
def create_status_file(self, extension, status, code, operation, message, overwrite=False):
# Create status file for specified extension. If overwrite is true, overwrite any existing status file. If
# false, create a status file only if it does not already exist.
_, status_path = self.get_status_file_path(extension)
if status_path is not None and not os.path.exists(status_path):
if status_path is not None and (overwrite or not os.path.exists(status_path)):
now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
status_contents = [
{
Expand Down
12 changes: 6 additions & 6 deletions tests/ga/test_extension.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3536,7 +3536,7 @@ def _create_policy_file(self, policy):
policy_file.write(policy)
policy_file.flush()

def _test_policy_case(self, policy, op, expected_status_code, expected_handler_status, expected_ext_count=1,
def _test_policy_case(self, policy, op, expected_status_code, expected_handler_status, expected_ext_count=0,
expected_status_msg=None):

with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
Expand Down Expand Up @@ -3630,13 +3630,13 @@ def test_should_fail_enable_if_dependent_extension_disallowed(self):

# OtherExampleHandlerLinux should be disallowed by policy, ExampleHandlerLinux should be skipped because
# dependent extension failed
self._assert_handler_status(protocol.report_vm_status, "NotReady", 1, "1.0.0",
expected_handler_name="OSTCExtensions.OtherExampleHandlerLinux",
self._assert_handler_status(protocol.report_vm_status, expected_status="NotReady", expected_ext_count=0,
version="1.0.0", expected_handler_name="OSTCExtensions.OtherExampleHandlerLinux",
expected_msg=("failed to run extension 'OSTCExtensions.OtherExampleHandlerLinux' "
"because it is not specified in the allowlist."))

self._assert_handler_status(protocol.report_vm_status, "NotReady", 1, "1.0.0",
expected_handler_name="OSTCExtensions.ExampleHandlerLinux",
self._assert_handler_status(protocol.report_vm_status, expected_status="NotReady", expected_ext_count=0,
version="1.0.0", expected_handler_name="OSTCExtensions.ExampleHandlerLinux",
expected_msg="Skipping processing of extensions since execution of dependent "
"extension OSTCExtensions.OtherExampleHandlerLinux failed")

Expand Down Expand Up @@ -3669,7 +3669,7 @@ def test_enable_should_succeed_if_extension_allowed(self):
]
for policy in policy_cases:
self._test_policy_case(policy=policy, op=ExtensionRequestedState.Enabled, expected_status_code=0,
expected_handler_status='Ready')
expected_handler_status='Ready', expected_ext_count=1)

def test_uninstall_should_succeed_if_extension_allowed(self):
policy_cases = [
Expand Down
Loading

0 comments on commit ba3869c

Please sign in to comment.