Repository: Azure/WALinuxAgent
Pull request #1122 (2018-04-19 04:44:14)

overwrite if an archived file exists
There were two failures cases that need to be addressed.

 1. If purge() is called before archive() the history directory may not exist yet.  The first call should create the directory if it does not exist.
 1. If the same file is archived it should be overwritten.  The same file was being overwritten because the timestamp was defined at class instantiation time instead of at flush time.


Patch:
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index 5df35a98..d7709b3e 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -16,6 +16,8 @@
 #
 # Requires Python 2.6+ and Openssl 1.0+
 
+from datetime import datetime
+
 import json
 import os
 import random
@@ -719,7 +721,7 @@ class WireClient(object):
                             # Goalstate is not updated.
                             return
 
-                self.goal_state_flusher.flush()
+                self.goal_state_flusher.flush(datetime.utcnow())
 
                 self.goal_state = goal_state
                 file_name = GOAL_STATE_FILE_NAME.format(goal_state.incarnation)
@@ -1129,7 +1131,8 @@ class WireClient(object):
 
                         host = self.get_host_plugin()
                         uri, headers = host.get_artifact_request(blob)
-                        profile = self.fetch(uri, headers, use_proxy=False)
+                        config = self.fetch(uri, headers, use_proxy=False)
+                        profile = self.decode_config(config)
 
                     if not textutil.is_str_none_or_whitespace(profile):
                         logger.verbose("Artifacts profile downloaded")
diff --git a/azurelinuxagent/common/utils/archive.py b/azurelinuxagent/common/utils/archive.py
index a167cff2..fd359d88 100644
--- a/azurelinuxagent/common/utils/archive.py
+++ b/azurelinuxagent/common/utils/archive.py
@@ -1,15 +1,14 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the Apache License.
-
 import os
 import re
 import shutil
 import zipfile
 
-from datetime import datetime
-
 from azurelinuxagent.common.utils import fileutil
 
+import azurelinuxagent.common.logger as logger
+
 
 """
 archive.py
@@ -34,6 +33,8 @@ is /var/lib/waagent.
 The timestamp is an ISO8601 formatted value.
 """
 
+ARCHIVE_DIRECTORY_NAME = 'history'
+
 MAX_ARCHIVED_STATES = 50
 
 CACHE_PATTERNS = [
@@ -50,20 +51,21 @@ ARCHIVE_PATTERNS_ZIP       = re.compile('^\d{4}\-\d{2}\-\d{2}T\d{2}:\d{2}:\d{2}\
 
 class StateFlusher(object):
     def __init__(self, lib_dir):
-        self._timestamp = datetime.now()
         self._source = lib_dir
-        self._target = os.path.join(self._source, 'history', self._timestamp.isoformat())
 
-    def flush(self):
+    def flush(self, timestamp):
         files = self._get_files_to_archive()
         if len(files) == 0:
             return
 
-        if self._mkdir():
-            self._archive(files)
+        if self._mkdir(timestamp):
+            self._archive(files, timestamp)
         else:
             self._purge(files)
 
+    def history_dir(self, timestamp):
+        return os.path.join(self._source, ARCHIVE_DIRECTORY_NAME, timestamp.isoformat())
+
     def _get_files_to_archive(self):
         files = []
         for f in os.listdir(self._source):
@@ -76,19 +78,23 @@ class StateFlusher(object):
 
         return files
 
-    def _archive(self, files):
+    def _archive(self, files, timestamp):
         for f in files:
-            shutil.move(f, self._target)
+            dst = os.path.join(self.history_dir(timestamp), os.path.basename(f))
+            shutil.move(f, dst)
 
     def _purge(self, files):
         for f in files:
             os.remove(f)
 
-    def _mkdir(self):
+    def _mkdir(self, timestamp):
+        d = self.history_dir(timestamp)
+
         try:
-            fileutil.mkdir(self._target, mode=0o700)
+            fileutil.mkdir(d, mode=0o700)
             return True
-        except:
+        except IOError as e:
+            logger.error("{0} : {1}".format(d, e.strerror))
             return False
 
 
@@ -162,7 +168,13 @@ class StateDirectory(State):
 
 class StateArchiver(object):
     def __init__(self, lib_dir):
-        self._source = os.path.join(lib_dir, 'history')
+        self._source = os.path.join(lib_dir, ARCHIVE_DIRECTORY_NAME)
+
+        if not os.path.isdir(self._source):
+            try:
+                fileutil.mkdir(self._source, mode=0o700)
+            except IOError as e:
+                logger.error("{0} : {1}", self._source, e.strerror)
 
     def purge(self):
         """


Test patch:
diff --git a/tests/common/test_event.py b/tests/common/test_event.py
index 7ab4b3a1..8ec6f082 100644
--- a/tests/common/test_event.py
+++ b/tests/common/test_event.py
@@ -108,7 +108,7 @@ class TestEvent(AgentTestCase):
         event.__event_logger__.reset_periodic()
 
         event.add_periodic(logger.EVERY_DAY, "FauxEvent")
-        self.assertEqual(1, mock_event.call_count)
+        mock_event.assert_called_once()
 
     @patch('azurelinuxagent.common.event.EventLogger.add_event')
     def test_periodic_does_not_emit_if_previously_sent(self, mock_event):
diff --git a/tests/common/test_logger.py b/tests/common/test_logger.py
index e28d385e..6a7e3796 100644
--- a/tests/common/test_logger.py
+++ b/tests/common/test_logger.py
@@ -34,7 +34,7 @@ class TestLogger(AgentTestCase):
         logger.reset_periodic()
 
         logger.periodic(logger.EVERY_DAY, _MSG, *_DATA)
-        self.assertEqual(1, mock_info.call_count)
+        mock_info.assert_called_once()
 
     @patch('azurelinuxagent.common.logger.Logger.info')
     def test_periodic_does_not_emit_if_previously_sent(self, mock_info):
diff --git a/tests/distro/test_resourceDisk.py b/tests/distro/test_resourceDisk.py
index d2ce6e16..6d96b0b1 100644
--- a/tests/distro/test_resourceDisk.py
+++ b/tests/distro/test_resourceDisk.py
@@ -78,7 +78,7 @@ class TestResourceDisk(AgentTestCase):
             # assert
             if sys.version_info >= (3,3):
                 with patch("os.posix_fallocate") as posix_fallocate:
-                    self.assertEqual(0, posix_fallocate.call_count)
+                    assert posix_fallocate.assert_not_called()
 
             assert run_patch.call_count == 1
             assert "dd if" in run_patch.call_args_list[0][0][0]
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 276d57e7..25bb5080 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -571,13 +571,13 @@ class TestExtension(AgentTestCase):
         conf.get_enable_overprovisioning = Mock(return_value=False)
         with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
             exthandlers_handler.handle_ext_handlers()
-            self.assertEqual(1, patch_handle_ext_handler.call_count)
+            patch_handle_ext_handler.assert_called()
 
         # enable extension handling blocking
         conf.get_enable_overprovisioning = Mock(return_value=True)
         with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
             exthandlers_handler.handle_ext_handlers()
-            self.assertEqual(0, patch_handle_ext_handler.call_count)
+            patch_handle_ext_handler.assert_not_called()
 
 
     def test_handle_ext_handlers_on_hold_false(self, *args):
@@ -596,13 +596,13 @@ class TestExtension(AgentTestCase):
         protocol.get_artifacts_profile = Mock(return_value=mock_in_vm_artifacts_profile)
         with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
             exthandlers_handler.handle_ext_handlers()
-            self.assertEqual(1, patch_handle_ext_handler.call_count)
+            patch_handle_ext_handler.assert_called_once()
 
         #Test when in_vm_artifacts_profile is not available
         protocol.get_artifacts_profile = Mock(return_value=None)
         with patch.object(ExtHandlersHandler, 'handle_ext_handler') as patch_handle_ext_handler:
             exthandlers_handler.handle_ext_handlers()
-            self.assertEqual(1, patch_handle_ext_handler.call_count)
+            patch_handle_ext_handler.assert_called_once()
 
     def _assert_ext_status(self, report_ext_status, expected_status,
                            expected_seq_no):
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 81d85f11..947dd0de 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -884,7 +884,7 @@ class TestUpdate(UpdateTestCase):
         mock_get_host.return_value = "faux host"
         host = self.update_handler._get_host_plugin(protocol=protocol)
         print("mock_get_host call cound={0}".format(mock_get_host.call_count))
-        self.assertEqual(1, mock_get_host.call_count)
+        mock_get_host.assert_called_once()
         self.assertEqual("faux host", host)
 
     @patch('azurelinuxagent.common.protocol.wire.WireClient.get_host_plugin')
diff --git a/tests/pa/test_provision.py b/tests/pa/test_provision.py
index 4ee8d469..6ec380cc 100644
--- a/tests/pa/test_provision.py
+++ b/tests/pa/test_provision.py
@@ -65,9 +65,9 @@ class TestProvision(AgentTestCase):
 
         ph.run()
 
-        self.assertEqual(0, ph.is_provisioned.call_count)
-        self.assertEqual(1, ph.report_ready.call_count)
-        self.assertEqual(1, ph.write_provisioned.call_count)
+        ph.is_provisioned.assert_not_called()
+        ph.report_ready.assert_called_once()
+        ph.write_provisioned.assert_called_once()
 
     @patch('os.path.isfile', return_value=False)
     def test_is_provisioned_not_provisioned(self, mock_isfile):
@@ -90,8 +90,8 @@ class TestProvision(AgentTestCase):
         mock_deprovision.return_value = deprovision_handler
 
         self.assertTrue(ph.is_provisioned())
-        self.assertEqual(1, ph.osutil.is_current_instance_id.call_count)
-        self.assertEqual(0, deprovision_handler.run_changed_unique_id.call_count)
+        ph.osutil.is_current_instance_id.assert_called_once()
+        deprovision_handler.run_changed_unique_id.assert_not_called()
 
     @patch('os.path.isfile', return_value=True)
     @patch('azurelinuxagent.common.utils.fileutil.read_file',
@@ -110,8 +110,8 @@ class TestProvision(AgentTestCase):
         mock_deprovision.return_value = deprovision_handler
 
         self.assertTrue(ph.is_provisioned())
-        self.assertEqual(1, ph.osutil.is_current_instance_id.call_count)
-        self.assertEqual(1, deprovision_handler.run_changed_unique_id.call_count)
+        ph.osutil.is_current_instance_id.assert_called_once()
+        deprovision_handler.run_changed_unique_id.assert_called_once()
 
     @distros()
     @patch('azurelinuxagent.common.osutil.default.DefaultOSUtil.get_instance_id',
diff --git a/tests/protocol/test_wire.py b/tests/protocol/test_wire.py
index 99dc5596..34b82862 100644
--- a/tests/protocol/test_wire.py
+++ b/tests/protocol/test_wire.py
@@ -153,7 +153,7 @@ class TestWireProtocol(AgentTestCase):
             host_plugin = wire_protocol_client.get_host_plugin()
             self.assertEqual(goal_state.container_id, host_plugin.container_id)
             self.assertEqual(goal_state.role_config_name, host_plugin.role_config_name)
-            self.assertEqual(1, patch_get_goal_state.call_count)
+            patch_get_goal_state.assert_called_once()
 
     def test_download_ext_handler_pkg_fallback(self, *args):
         ext_uri = 'extension_uri'
@@ -218,7 +218,7 @@ class TestWireProtocol(AgentTestCase):
                     wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
                     wire_protocol_client.upload_status_blob()
                     patch_default_upload.assert_called_once_with(testurl)
-                    self.assertEqual(1, wire_protocol_client.get_goal_state.call_count)
+                    wire_protocol_client.get_goal_state.assert_called_once()
                     patch_http.assert_called_once_with(testurl, wire_protocol_client.status_blob)
                     self.assertTrue(HostPluginProtocol.is_default_channel())
                     HostPluginProtocol.set_default_channel(False)
@@ -255,8 +255,8 @@ class TestWireProtocol(AgentTestCase):
             with patch.object(WireClient, "report_status_event") as mock_event:
                 wire_protocol_client.upload_status_blob()
 
-                self.assertEqual(1, mock_prepare.call_count)
-                self.assertEqual(1, mock_event.call_count)
+                mock_prepare.assert_called_once()
+                mock_event.assert_called_once()
 
     def test_get_in_vm_artifacts_profile_blob_not_available(self, *args):
         wire_protocol_client = WireProtocol(wireserver_url).client
@@ -342,7 +342,7 @@ class TestWireProtocol(AgentTestCase):
         wire_protocol_client.ext_conf.artifacts_profile_blob = testurl
         goal_state = GoalState(WireProtocolData(DATA_FILE).goal_state)
         wire_protocol_client.get_goal_state = Mock(return_value=goal_state)
-        wire_protocol_client.fetch = Mock(side_effect=[None, '{"onHold": "true"}'])
+        wire_protocol_client.fetch = Mock(side_effect=[None, '{"onHold": "true"}'.encode('utf-8')])
         with patch.object(HostPluginProtocol,
                           "get_artifact_request",
                           return_value=['dummy_url', {}]) as artifact_request:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 0f71259a..ab04b9c2 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -117,7 +117,7 @@ class TestAgent(AgentTestCase):
         agent.daemon()
 
         mock_daemon.run.assert_called_once_with(child_args=None)
-        self.assertEqual(1, mock_load.call_count)
+        mock_load.assert_called_once()
 
     @patch("azurelinuxagent.daemon.get_daemon_handler")
     @patch("azurelinuxagent.common.conf.load_conf_from_file")
@@ -131,7 +131,7 @@ class TestAgent(AgentTestCase):
         agent.daemon()
 
         mock_daemon.run.assert_called_once_with(child_args="-configuration-path:/foo/bar.conf")
-        self.assertEqual(1, mock_load.call_count)
+        mock_load.assert_called_once()
 
     @patch("azurelinuxagent.common.conf.get_ext_log_dir")
     def test_agent_ensures_extension_log_directory(self, mock_dir):
@@ -156,7 +156,7 @@ class TestAgent(AgentTestCase):
                     conf_file_path=os.path.join(data_dir, "test_waagent.conf"))
         self.assertTrue(os.path.isfile(ext_log_dir))
         self.assertFalse(os.path.isdir(ext_log_dir))
-        self.assertEqual(1, mock_log.call_count)
+        mock_log.assert_called_once()
 
     def test_agent_get_configuration(self):
         Agent(False, conf_file_path=os.path.join(data_dir, "test_waagent.conf"))
diff --git a/tests/utils/test_archive.py b/tests/utils/test_archive.py
index 1cd251e4..f0564c88 100644
--- a/tests/utils/test_archive.py
+++ b/tests/utils/test_archive.py
@@ -1,11 +1,9 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the Apache License.
+from datetime import datetime, timedelta
 
 import zipfile
 
-from datetime import datetime, timedelta
-
-from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.common.utils.archive import StateFlusher, StateArchiver, MAX_ARCHIVED_STATES
 from tests.tools import *
 
@@ -57,7 +55,7 @@ class TestArchive(AgentTestCase):
             self._write_file(f)
 
         test_subject = StateFlusher(self.tmp_dir)
-        test_subject.flush()
+        test_subject.flush(datetime.utcnow())
 
         self.assertTrue(os.path.exists(self.history_dir))
         self.assertTrue(os.path.isdir(self.history_dir))
@@ -67,7 +65,7 @@ class TestArchive(AgentTestCase):
 
         self.assertIsIso8601(timestamp_dirs[0])
         ts = self.parse_isoformat(timestamp_dirs[0])
-        self.assertDateTimeCloseTo(ts, datetime.now(), timedelta(seconds=30))
+        self.assertDateTimeCloseTo(ts, datetime.utcnow(), timedelta(seconds=30))
 
         for f in temp_files:
             history_path = os.path.join(self.history_dir, timestamp_dirs[0], f)
@@ -92,7 +90,7 @@ class TestArchive(AgentTestCase):
             self._write_file(f)
 
         flusher = StateFlusher(self.tmp_dir)
-        flusher.flush()
+        flusher.flush(datetime.utcnow())
 
         test_subject = StateArchiver(self.tmp_dir)
         test_subject.archive()
@@ -105,7 +103,7 @@ class TestArchive(AgentTestCase):
 
         self.assertIsIso8601(ts_s)
         ts = self.parse_isoformat(ts_s)
-        self.assertDateTimeCloseTo(ts, datetime.now(), timedelta(seconds=30))
+        self.assertDateTimeCloseTo(ts, datetime.utcnow(), timedelta(seconds=30))
 
         zip_full = os.path.join(self.history_dir, zip_fn)
         self.assertZipContains(zip_full, temp_files)
@@ -155,6 +153,69 @@ class TestArchive(AgentTestCase):
                 fn = "{0}.zip".format(ts)
             self.assertTrue(fn in archived_entries, "'{0}' is not in the list of unpurged entires".format(fn))
 
+    def test_archive03(self):
+        """
+        If the StateFlusher has to flush the same file, it should
+        overwrite the existing one.
+        """
+        temp_files = [
+            'Prod.0.manifest.xml',
+            'Prod.0.agentsManifest',
+            'Microsoft.Azure.Extensions.CustomScript.0.xml'
+        ]
+
+        def _write_goal_state_files(temp_files):
+            for f in temp_files:
+                self._write_file(f)
+
+        timestamp = datetime.utcnow()
+
+        _write_goal_state_files(temp_files)
+        test_subject = StateFlusher(self.tmp_dir)
+        test_subject.flush(timestamp)
+
+        # Ensure that the file contets are what we expect after a flush.
+        # The file contents are simply the name of the file.
+        fn = os.path.join(self.history_dir, timestamp.isoformat(), 'Prod.0.manifest.xml')
+        self.assertEqual('Prod.0.manifest.xml', fileutil.read_file(fn))
+
+        # Modify the contents of the file on disk, and re-flush().
+        fileutil.write_file(fn, "--this-has-been-changed--")
+        self.assertEqual("--this-has-been-changed--", fileutil.read_file(fn))
+
+        # re-write all of the same files, and flush again. .flush()
+        # should overwrite the existing ones
+
+        _write_goal_state_files(temp_files)
+        test_subject.flush(timestamp)
+
+        # The contents of the file were overwritten as a result of the flush.
+        self.assertEqual('Prod.0.manifest.xml', fileutil.read_file(fn))
+
+        self.assertTrue(os.path.exists(self.history_dir))
+        self.assertTrue(os.path.isdir(self.history_dir))
+
+        timestamp_dirs = os.listdir(self.history_dir)
+        self.assertEqual(1, len(timestamp_dirs))
+
+        self.assertIsIso8601(timestamp_dirs[0])
+        ts = self.parse_isoformat(timestamp_dirs[0])
+        self.assertDateTimeCloseTo(ts, datetime.utcnow(), timedelta(seconds=30))
+
+        for f in temp_files:
+            history_path = os.path.join(self.history_dir, timestamp_dirs[0], f)
+            msg = "expected the temp file {0} to exist".format(history_path)
+            self.assertTrue(os.path.exists(history_path), msg)
+
+    def test_archive04(self):
+        """
+        The archive directory is created if it does not exist.
+
+        This failure was caught when .purge() was called before .archive().
+        """
+        test_subject = StateArchiver(os.path.join(self.tmp_dir, 'does-not-exist'))
+        test_subject.purge()
+
     def parse_isoformat(self, s):
         return datetime.strptime(s, '%Y-%m-%dT%H:%M:%S.%f')
 
diff --git a/tests/utils/test_rest_util.py b/tests/utils/test_rest_util.py
index a864884a..e5aee4e8 100644
--- a/tests/utils/test_rest_util.py
+++ b/tests/utils/test_rest_util.py
@@ -132,8 +132,8 @@ class TestHttpOperations(AgentTestCase):
         h, p = restutil._get_http_proxy()
         self.assertEqual("host", h)
         self.assertEqual(None, p)
-        self.assertEqual(1, mock_host.call_count)
-        self.assertEqual(1, mock_port.call_count)
+        mock_host.assert_called_once()
+        mock_port.assert_called_once()
 
     @patch('azurelinuxagent.common.conf.get_httpproxy_port')
     @patch('azurelinuxagent.common.conf.get_httpproxy_host')
@@ -143,8 +143,8 @@ class TestHttpOperations(AgentTestCase):
         h, p = restutil._get_http_proxy()
         self.assertEqual(None, h)
         self.assertEqual(None, p)
-        self.assertEqual(1, mock_host.call_count)
-        self.assertEqual(0, mock_port.call_count)
+        mock_host.assert_called_once()
+        mock_port.assert_not_called()
 
     @patch('azurelinuxagent.common.conf.get_httpproxy_host')
     def test_get_http_proxy_http_uses_httpproxy(self, mock_host):
@@ -197,7 +197,7 @@ class TestHttpOperations(AgentTestCase):
         mock_conn.request.assert_has_calls([
             call(method="GET", url="/bar", body=None, headers={'User-Agent': HTTP_USER_AGENT})
         ])
-        self.assertEqual(1, mock_conn.getresponse.call_count)
+        mock_conn.getresponse.assert_called_once()
         self.assertNotEquals(None, resp)
         self.assertEquals("TheResults", resp.read())
 
@@ -220,7 +220,7 @@ class TestHttpOperations(AgentTestCase):
         mock_conn.request.assert_has_calls([
             call(method="GET", url="/bar", body=None, headers={'User-Agent': HTTP_USER_AGENT})
         ])
-        self.assertEqual(1, mock_conn.getresponse.call_count)
+        mock_conn.getresponse.assert_called_once()
         self.assertNotEquals(None, resp)
         self.assertEquals("TheResults", resp.read())
 
@@ -244,7 +244,7 @@ class TestHttpOperations(AgentTestCase):
         mock_conn.request.assert_has_calls([
             call(method="GET", url="http://foo:80/bar", body=None, headers={'User-Agent': HTTP_USER_AGENT})
         ])
-        self.assertEqual(1, mock_conn.getresponse.call_count)
+        mock_conn.getresponse.assert_called_once()
         self.assertNotEquals(None, resp)
         self.assertEquals("TheResults", resp.read())
 
@@ -269,7 +269,7 @@ class TestHttpOperations(AgentTestCase):
         mock_conn.request.assert_has_calls([
             call(method="GET", url="https://foo:443/bar", body=None, headers={'User-Agent': HTTP_USER_AGENT})
         ])
-        self.assertEqual(1, mock_conn.getresponse.call_count)
+        mock_conn.getresponse.assert_called_once()
         self.assertNotEquals(None, resp)
         self.assertEquals("TheResults", resp.read())