Add update_file context manager for local files

The subprocess.update_file function provides an ability to update
content of a local file by iterating over lines of an original file and
forming a result content in a temporary file to replace the original
file in the end. This function is very useful to change configuration
files.

Change-Id: I433a5da67887b231400dd63131799019f45c277c

octane/tests/test_util_subprocess.py

@@ -0,0 +1,68 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import mock
+import pytest
+
+from octane.util import subprocess
+
+
+class _TestException(Exception):
+    pass
+
+
+@pytest.mark.parametrize(("exception", "reraise", "calls"), [
+    (None, False, [
+        mock.call.stat("/fake/filename"),
+        mock.call.chmod("/temp/filename", 0o640),
+        mock.call.chown("/temp/filename", 2, 3),
+        mock.call.rename("/fake/filename", "/fake/filename.bak"),
+        mock.call.rename("/temp/filename", "/fake/filename"),
+        mock.call.unlink("/fake/filename.bak"),
+    ]),
+    (subprocess.DontUpdateException, False, [
+        mock.call.unlink("/temp/filename"),
+    ]),
+    (_TestException, True, [
+        mock.call.unlink("/temp/filename"),
+    ]),
+])
+def test_update_file(mocker, mock_open, exception, reraise, calls):
+    mock_tempfile = mocker.patch("octane.util.tempfile.get_tempname")
+    mock_tempfile.return_value = "/temp/filename"
+
+    mock_old = mock.MagicMock()
+    mock_new = mock.MagicMock()
+
+    mock_open.side_effect = [mock_old, mock_new]
+
+    mock_os = mock.Mock()
+    os_methods = ["unlink", "stat", "chmod", "chown", "rename"]
+    for method in os_methods:
+        mocker.patch("os." + method, new=getattr(mock_os, method))
+
+    mock_os.stat.return_value.configure_mock(
+        st_mode=0o640,
+        st_uid=2,
+        st_gid=3,
+    )
+
+    if reraise:
+        with pytest.raises(exception):
+            with subprocess.update_file("/fake/filename"):
+                raise exception
+    else:
+        with subprocess.update_file("/fake/filename"):
+            if exception is not None:
+                raise exception
+
+    assert mock_os.mock_calls == calls

octane/util/maintenance.py

@@ -44,7 +44,7 @@ def disable_apis(env):
             with ssh.update_file(sftp, f) as (old, new):
                 contents = old.read()
                 if not mode_tcp_re.search(contents):
-                    raise ssh.DontUpdateException
+                    raise subprocess.DontUpdateException
                 new.write(contents)
                 if not contents.endswith('\n'):
                     new.write('\n')

octane/util/ssh.py

@@ -191,10 +191,6 @@ def sftp(node):
     return _get_sftp(node)
 
 
-class DontUpdateException(Exception):
-    pass
-
-
 @contextlib.contextmanager
 def update_file(sftp, filename):
     old = sftp.open(filename, 'r')
@@ -209,7 +205,7 @@ def update_file(sftp, filename):
     with contextlib.nested(old, new):
         try:
             yield old, new
-        except DontUpdateException:
+        except subprocess.DontUpdateException:
             sftp.unlink(temp_filename)
             return
         except Exception:

octane/util/subprocess.py

@@ -22,6 +22,8 @@ import re
 import subprocess
 import threading
 
+from octane.util import tempfile
+
 LOG = logging.getLogger(__name__)
 PIPE = subprocess.PIPE
 CalledProcessError = subprocess.CalledProcessError
@@ -202,3 +204,32 @@ def call(cmd, **kwargs):
 
 def call_output(cmd, **kwargs):
     return call(cmd, stdout=PIPE, **kwargs)[0]
+
+
+class DontUpdateException(Exception):
+    pass
+
+
+@contextlib.contextmanager
+def update_file(filename):
+    old = open(filename, 'r')
+    dirname = os.path.dirname(filename)
+    prefix = ".{0}.".format(os.path.basename(filename))
+    temp_filename = tempfile.get_tempname(dir=dirname, prefix=prefix)
+    new = open(temp_filename, 'w')
+    with contextlib.nested(old, new):
+        try:
+            yield old, new
+        except DontUpdateException:
+            os.unlink(temp_filename)
+            return
+        except Exception:
+            os.unlink(temp_filename)
+            raise
+        stat = os.stat(filename)
+        os.chmod(temp_filename, stat.st_mode)
+        os.chown(temp_filename, stat.st_uid, stat.st_gid)
+    bak_filename = filename + '.bak'
+    os.rename(filename, bak_filename)
+    os.rename(temp_filename, filename)
+    os.unlink(bak_filename)