From 1da0ef6e3c0be655a890f621e38b97b0820f4971 Mon Sep 17 00:00:00 2001 From: Norberto Arrieta Date: Tue, 19 Apr 2022 20:37:40 -0700 Subject: [PATCH] Updated tests in test_update.py to use mock_update_handler (#2552) Co-authored-by: narrieta --- tests/ga/mocks.py | 97 ++++++++++++++++-------- tests/ga/test_update.py | 158 ++++++++++++++++++++-------------------- 2 files changed, 146 insertions(+), 109 deletions(-) diff --git a/tests/ga/mocks.py b/tests/ga/mocks.py index 3b93bc8bc7..6fbc63d7da 100644 --- a/tests/ga/mocks.py +++ b/tests/ga/mocks.py @@ -1,4 +1,4 @@ -# Copyright 2018 Microsoft Corporation + # Copyright 2018 Microsoft Corporation # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,17 +25,29 @@ @contextlib.contextmanager -def mock_update_handler(protocol, iterations=1, on_new_iteration=lambda _: None, exthandlers_handler=None, remote_access_handler=None, enable_agent_updates=False): +def mock_update_handler(protocol, + iterations=1, + on_new_iteration=lambda _: None, + exthandlers_handler=None, + remote_access_handler=None, + autoupdate_enabled=False, + check_daemon_running=False, + start_background_threads=False, + check_background_threads=False + ): """ Creates a mock UpdateHandler that executes its main loop for the given 'iterations'. * If 'on_new_iteration' is given, it is invoked at the beginning of each iteration passing the iteration number as argument. * Network requests (e.g. requests for the goal state) are done using the given 'protocol'. * The mock UpdateHandler uses mock no-op ExtHandlersHandler and RemoteAccessHandler, unless they are given by 'exthandlers_handler' and 'remote_access_handler'. - * Agent updates are disabled, unless specified otherwise with 'enable_agent_updates'. - * Background threads (monitor, env, telemetry, etc) are not started. + * Agent updates are disabled, unless specified otherwise by 'autoupdate_enabled'. + * The check for the daemon is skipped, unless specified otherwise by 'check_daemon_running' + * Background threads (monitor, env, telemetry, etc) are not started, unless specified otherwise by 'start_background_threads' + * The check for background threads is skipped, unless specified otherwise by 'check_background_threads' * The UpdateHandler is augmented with these extra functions: - * get_exit_code() - returns the code passed to sys.exit() when the handler exits - * get_iterations() - returns the number of iterations executed by the main loop + * get_exit_code() - returns the code passed to sys.exit() when the handler exits + * get_iterations() - returns the number of iterations executed by the main loop + * get_iterations_completed() - returns the number of iterations of the main loop that completed execution (i.e. were not interrupted by an exception, return, etc) """ iteration_count = [0] @@ -53,34 +65,55 @@ def is_running(*args): # mock for property UpdateHandler.is_running, which cont if exthandlers_handler is None: exthandlers_handler = ExtHandlersHandler(protocol) + else: + exthandlers_handler.protocol = protocol if remote_access_handler is None: remote_access_handler = RemoteAccessHandler(protocol) - with patch("azurelinuxagent.ga.exthandlers.get_exthandlers_handler", return_value=exthandlers_handler): - with patch("azurelinuxagent.ga.remoteaccess.get_remote_access_handler", return_value=remote_access_handler): - with patch("azurelinuxagent.common.conf.get_autoupdate_enabled", return_value=enable_agent_updates): - with patch.object(UpdateHandler, "is_running", PropertyMock(side_effect=is_running)): - with patch.object(UpdateHandler, "_check_daemon_running"): - with patch.object(UpdateHandler, "_start_threads"): - with patch.object(UpdateHandler, "_check_threads_running"): - with patch('time.sleep', side_effect=lambda _: mock_sleep(0.001)): - with patch('sys.exit', side_effect=lambda _: 0) as mock_exit: - def get_exit_code(): - if mock_exit.call_count == 0: - raise Exception("The UpdateHandler did not exit") - if mock_exit.call_count != 1: - raise Exception("The UpdateHandler exited multiple times ({0})".format(mock_exit.call_count)) - args, _ = mock_exit.call_args - return args[0] - - def get_iterations(): - return iteration_count[0] - - update_handler = get_update_handler() - update_handler.protocol_util.get_protocol = Mock(return_value=protocol) - update_handler.get_exit_code = get_exit_code - update_handler.get_iterations = get_iterations - - yield update_handler + cleanup_functions = [] + + def patch_object(target, attribute): + p = patch.object(target, attribute) + p.start() + cleanup_functions.insert(0, p.stop) + + try: + with patch("azurelinuxagent.ga.exthandlers.get_exthandlers_handler", return_value=exthandlers_handler): + with patch("azurelinuxagent.ga.remoteaccess.get_remote_access_handler", return_value=remote_access_handler): + with patch("azurelinuxagent.common.conf.get_autoupdate_enabled", return_value=autoupdate_enabled): + with patch.object(UpdateHandler, "is_running", PropertyMock(side_effect=is_running)): + with patch('azurelinuxagent.ga.update.time.sleep', side_effect=lambda _: mock_sleep(0.001)) as sleep: + with patch('sys.exit', side_effect=lambda _: 0) as mock_exit: + if not check_daemon_running: + patch_object(UpdateHandler, "_check_daemon_running") + if not start_background_threads: + patch_object(UpdateHandler, "_start_threads") + if not check_background_threads: + patch_object(UpdateHandler, "_check_threads_running") + + def get_exit_code(): + if mock_exit.call_count == 0: + raise Exception("The UpdateHandler did not exit") + if mock_exit.call_count != 1: + raise Exception("The UpdateHandler exited multiple times ({0})".format(mock_exit.call_count)) + args, _ = mock_exit.call_args + return args[0] + + def get_iterations(): + return iteration_count[0] + + def get_iterations_completed(): + return sleep.call_count + + update_handler = get_update_handler() + update_handler.protocol_util.get_protocol = Mock(return_value=protocol) + update_handler.get_exit_code = get_exit_code + update_handler.get_iterations = get_iterations + update_handler.get_iterations_completed = get_iterations_completed + + yield update_handler + finally: + for f in cleanup_functions: + f() diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py index 8cc4d59828..664c78e205 100644 --- a/tests/ga/test_update.py +++ b/tests/ga/test_update.py @@ -112,10 +112,10 @@ def _get_update_handler(iterations=1, test_data=None, protocol=None): with patch.object(HostPluginProtocol, "is_default_channel", False): if protocol is None: with mock_wire_protocol(test_data) as mock_protocol: - with mock_update_handler(mock_protocol, iterations=iterations, enable_agent_updates=True) as update_handler: + with mock_update_handler(mock_protocol, iterations=iterations, autoupdate_enabled=True) as update_handler: yield update_handler, mock_protocol else: - with mock_update_handler(protocol, iterations=iterations, enable_agent_updates=True) as update_handler: + with mock_update_handler(protocol, iterations=iterations, autoupdate_enabled=True) as update_handler: yield update_handler, protocol @@ -1294,77 +1294,6 @@ def test_get_latest_agent_should_return_latest_agent_even_on_bad_error_json(self latest_agent = self.update_handler.get_latest_agent_greater_than_daemon() self.assertEqual(latest_agent.version, dst_ver, "Latest agent version is invalid") - def _test_run(self, invocations=1, calls=1, enable_updates=False, sleep_interval=(6,)): - conf.get_autoupdate_enabled = Mock(return_value=enable_updates) - - def iterator(*_, **__): - iterator.count += 1 - if iterator.count <= invocations: - return True - return False - iterator.count = 0 - - fileutil.write_file(conf.get_agent_pid_file_path(), ustr(42)) - - with patch('azurelinuxagent.ga.exthandlers.get_exthandlers_handler') as mock_handler: - mock_handler.run_ext_handlers = Mock() - with patch('azurelinuxagent.ga.update.get_monitor_handler') as mock_monitor: - with patch.object(UpdateHandler, 'is_running') as mock_is_running: - mock_is_running.__get__ = Mock(side_effect=iterator) - with patch('azurelinuxagent.ga.remoteaccess.get_remote_access_handler') as mock_ra_handler: - with patch('azurelinuxagent.ga.update.get_env_handler') as mock_env: - with patch('azurelinuxagent.ga.update.get_collect_logs_handler') as mock_collect_logs: - with patch('azurelinuxagent.ga.update.get_send_telemetry_events_handler') as mock_telemetry_send_events: - with patch('azurelinuxagent.ga.update.get_collect_telemetry_events_handler') as mock_event_collector: - with patch('azurelinuxagent.ga.update.initialize_event_logger_vminfo_common_parameters'): - with patch('azurelinuxagent.ga.update.is_log_collection_allowed', return_value=True): - with patch.object(self.update_handler, "_processing_new_extensions_goal_state", return_value=True): - with patch('time.sleep') as sleep_mock: - with patch('sys.exit') as mock_exit: - if isinstance(os.getppid, MagicMock): - self.update_handler.run() - else: - with patch('os.getppid', return_value=42): - self.update_handler.run() - - self.assertEqual(1, mock_handler.call_count) - self.assertEqual(calls, len([c for c in [call[0] for call in mock_handler.return_value.method_calls] if c == 'run'])) - self.assertEqual(1, mock_ra_handler.call_count) - self.assertEqual(calls, len(mock_ra_handler.return_value.method_calls)) - if calls > 0: - self.assertEqual(sleep_interval, sleep_mock.call_args[0]) - self.assertEqual(1, mock_monitor.call_count) - self.assertEqual(1, mock_env.call_count) - self.assertEqual(1, mock_collect_logs.call_count) - self.assertEqual(1, mock_telemetry_send_events.call_count) - self.assertEqual(1, mock_event_collector.call_count) - self.assertEqual(1, mock_exit.call_count) - - def test_run(self): - self._test_run() - - def test_run_stops_if_update_available(self): - self.update_handler._download_agent_if_upgrade_available = Mock(return_value=True) - self._test_run(invocations=0, calls=0, enable_updates=True) - - def test_run_stops_if_orphaned(self): - with patch('os.getppid', return_value=1): - self._test_run(invocations=0, calls=0, enable_updates=True) - - def test_run_clears_sentinel_on_successful_exit(self): - self._test_run() - self.assertFalse(os.path.isfile(self.update_handler._sentinel_file_path())) - - def test_run_leaves_sentinel_on_unsuccessful_exit(self): - self.update_handler._download_agent_if_upgrade_available = Mock(side_effect=Exception) - self._test_run(invocations=1, calls=0, enable_updates=True) - self.assertTrue(os.path.isfile(self.update_handler._sentinel_file_path())) - - def test_run_emits_restart_event(self): - self.update_handler._emit_restart_event = Mock() - self._test_run() - self.assertEqual(1, self.update_handler._emit_restart_event.call_count) - def test_set_agents_sets_agents(self): self.prepare_agents() @@ -1510,15 +1439,19 @@ def test_write_pid_file_ignores_exceptions(self): self.assertEqual(0, len(pid_files)) self.assertEqual(None, pid_file) - @patch('azurelinuxagent.common.conf.get_extensions_enabled', return_value=False) - def test_update_happens_when_extensions_disabled(self, _): + def test_update_happens_when_extensions_disabled(self): """ Although the extension enabled config will not get checked before an update is found, this test attempts to ensure that behavior never changes. """ - self.update_handler._download_agent_if_upgrade_available = Mock(return_value=True) - self._test_run(invocations=0, calls=0, enable_updates=True, sleep_interval=(300,)) + with patch('azurelinuxagent.common.conf.get_extensions_enabled', return_value=False): + with patch('azurelinuxagent.ga.update.UpdateHandler._download_agent_if_upgrade_available', return_value=True) as download_agent: + with mock_wire_protocol(DATA_FILE) as protocol: + with mock_update_handler(protocol, autoupdate_enabled=True) as update_handler: + update_handler.run() + + self.assertEqual(1, download_agent.call_count, "Agent update did not execute (no attempts to download the agent") @staticmethod def _get_test_ext_handler_instance(protocol, name="OSTCExtensions.ExampleHandlerLinux", version="1.0.0"): @@ -1897,6 +1830,77 @@ def test_it_should_reset_legacy_blacklisted_agents_on_process_start(self): self.assertFalse(agent.is_blacklisted, "Legacy Agent should not be blacklisted") +class UpdateHandlerRunTestCase(AgentTestCase): + def _test_run(self, autoupdate_enabled=False, check_daemon_running=False, expected_exit_code=0, emit_restart_event=None): + fileutil.write_file(conf.get_agent_pid_file_path(), ustr(42)) + + with patch('azurelinuxagent.ga.update.get_monitor_handler') as mock_monitor: + with patch('azurelinuxagent.ga.remoteaccess.get_remote_access_handler') as mock_ra_handler: + with patch('azurelinuxagent.ga.update.get_env_handler') as mock_env: + with patch('azurelinuxagent.ga.update.get_collect_logs_handler') as mock_collect_logs: + with patch('azurelinuxagent.ga.update.get_send_telemetry_events_handler') as mock_telemetry_send_events: + with patch('azurelinuxagent.ga.update.get_collect_telemetry_events_handler') as mock_event_collector: + with patch('azurelinuxagent.ga.update.initialize_event_logger_vminfo_common_parameters'): + with patch('azurelinuxagent.ga.update.is_log_collection_allowed', return_value=True): + with mock_wire_protocol(DATA_FILE) as protocol: + mock_exthandlers_handler = Mock() + with mock_update_handler( + protocol, + exthandlers_handler=mock_exthandlers_handler, + remote_access_handler=mock_ra_handler, + autoupdate_enabled=autoupdate_enabled, + check_daemon_running=check_daemon_running + ) as update_handler: + + if emit_restart_event is not None: + update_handler._emit_restart_event = emit_restart_event + + if isinstance(os.getppid, MagicMock): + update_handler.run() + else: + with patch('os.getppid', return_value=42): + update_handler.run() + + self.assertEqual(1, mock_monitor.call_count) + self.assertEqual(1, mock_env.call_count) + self.assertEqual(1, mock_collect_logs.call_count) + self.assertEqual(1, mock_telemetry_send_events.call_count) + self.assertEqual(1, mock_event_collector.call_count) + self.assertEqual(expected_exit_code, update_handler.get_exit_code()) + + if update_handler.get_iterations_completed() > 0: # some test cases exit before executing extensions or remote access + self.assertEqual(1, mock_exthandlers_handler.run.call_count) + self.assertEqual(1, mock_ra_handler.run.call_count) + + return update_handler + + def test_run(self): + self._test_run() + + def test_run_stops_if_update_available(self): + with patch('azurelinuxagent.ga.update.UpdateHandler._download_agent_if_upgrade_available', return_value=True): + update_handler = self._test_run(autoupdate_enabled=True) + self.assertEqual(0, update_handler.get_iterations_completed()) + + def test_run_stops_if_orphaned(self): + with patch('os.getppid', return_value=1): + update_handler = self._test_run(check_daemon_running=True) + self.assertEqual(0, update_handler.get_iterations_completed()) + + def test_run_clears_sentinel_on_successful_exit(self): + update_handler = self._test_run() + self.assertFalse(os.path.isfile(update_handler._sentinel_file_path())) + + def test_run_leaves_sentinel_on_unsuccessful_exit(self): + with patch('azurelinuxagent.ga.update.UpdateHandler._download_agent_if_upgrade_available', side_effect=Exception): + update_handler = self._test_run(autoupdate_enabled=True,expected_exit_code=1) + self.assertTrue(os.path.isfile(update_handler._sentinel_file_path())) + + def test_run_emits_restart_event(self): + update_handler = self._test_run(emit_restart_event=Mock()) + self.assertEqual(1, update_handler._emit_restart_event.call_count) + + class TestAgentUpgrade(UpdateTestCase): @contextlib.contextmanager