Unpack bootstrap img in context manager

Change-Id: I386ea151c5675f93da4359fadd8adaa6ba890e67

octane/commands/prepare.py

@@ -11,12 +11,11 @@
 # under the License.
 
 import os.path
-import shutil
-import tempfile
 
 from cliff import command as cmd
 
 from octane import magic_consts
+from octane.util import archivate
 from octane.util import docker
 from octane.util import subprocess
 
@@ -55,19 +54,9 @@ def revert_initramfs():
 
 
 def patch_initramfs():
-    backup = magic_consts.BOOTSTRAP_INITRAMFS + '.bkup'
-    chroot = tempfile.mkdtemp()
-    try:
-        os.rename(magic_consts.BOOTSTRAP_INITRAMFS, backup)
-        subprocess.call("gunzip -c {0} | cpio -id".format(backup),
-                        shell=True, cwd=chroot)
+    with archivate.update_cpio(magic_consts.BOOTSTRAP_INITRAMFS) as chroot:
         patch_fuel_agent(chroot)
-        with open(magic_consts.BOOTSTRAP_INITRAMFS, "wb") as f:
-            subprocess.call("find | grep -v '^\.$' | cpio --format newc -o"
-                            " | gzip -c", shell=True, stdout=f, cwd=chroot)
-        docker.run_in_container("cobbler", ["cobbler", "sync"])
-    finally:
-        shutil.rmtree(chroot)
+    docker.run_in_container("cobbler", ["cobbler", "sync"])
 
 
 def patch_fuel_agent(chroot):

octane/tests/test_archive.py

@@ -82,3 +82,75 @@ def test_archivate_container_cmd_output(mocker):
     test_archive.addfile.assert_called_once_with(
         tar_info, io_mock.return_value)
     assert io_mock.return_value.getvalue() == output_data
+
+
+def test_update_cpio_with_exception(mocker, mock_subprocess):
+    mocker.patch("tempfile.mkdtemp", return_value="tmp_dir_name")
+    shutil_mocker = mocker.patch("shutil.rmtree")
+
+    class TestException(Exception):
+        pass
+
+    test_msg = "Fake exception"
+    with pytest.raises(TestException) as exc_info:
+        with archivate.update_cpio("fake_img_path.img"):
+            raise TestException(test_msg)
+    assert test_msg == exc_info.value.message
+    shutil_mocker.assert_called_once_with("tmp_dir_name")
+    assert [
+        mock.call(
+            ["gunzip", "-c", "fake_img_path.img"],
+            stdout=subprocess.PIPE),
+        mock.call(
+            ["cpio", "-id"],
+            stdin=mock_subprocess.return_value.__enter__.return_value.stdout,
+            cwd="tmp_dir_name"),
+    ] == mock_subprocess.call_args_list
+
+
+def test_update_cpio(mocker, mock_subprocess):
+    tmp_dir_name = "tmp_dir_name"
+    mocker.patch("tempfile.mkdtemp", return_value=tmp_dir_name)
+    tmp_file = mocker.patch(
+        "tempfile.NamedTemporaryFile"
+    ).return_value.__enter__.return_value
+    shutil_rm_mocker = mocker.patch("shutil.rmtree")
+    shutil_mv_mocker = mocker.patch("shutil.move")
+    mock_subprocess.return_value.__enter__.return_value = \
+        mock_subprocess.return_value
+    walk_mocker = mocker.patch(
+        "os.walk",
+        return_value=[(os.path.join(tmp_dir_name, "path"),
+                       ["dir_1", "dir_2"],
+                       ["file_1", "file_2"])])
+    with archivate.update_cpio("fake_img_path.img") as tmp_dir_name_context:
+        pass
+    assert tmp_dir_name == tmp_dir_name_context
+    walk_mocker.assert_called_once_with(tmp_dir_name)
+    shutil_rm_mocker.assert_called_once_with(tmp_dir_name)
+    shutil_mv_mocker.assert_called_once_with(
+        tmp_file.name, "fake_img_path.img")
+    assert [
+        mock.call("path/{0}\n".format(i)) for i in
+        ["dir_1", "dir_2", "file_1", "file_2"]
+    ] == mock_subprocess.return_value.stdin.write.call_args_list
+    assert [
+        mock.call(
+            ["gunzip", "-c", "fake_img_path.img"],
+            stdout=subprocess.PIPE),
+        mock.call(
+            ["cpio", "-id"],
+            stdin=mock_subprocess.return_value.stdout,
+            cwd=tmp_dir_name),
+        mock.call(
+            ["cpio", "--format", "newc", "-o"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            cwd=tmp_dir_name),
+        mock.call(
+            ["gzip", "-c"],
+            stdin=mock_subprocess.return_value.stdout,
+            stdout=tmp_file,
+            cwd=tmp_dir_name),
+    ] == mock_subprocess.call_args_list
+    assert not tmp_file.delete

octane/util/archivate.py

@@ -10,9 +10,13 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import contextlib
 import io
+import itertools
 import os
+import shutil
 import tarfile
+import tempfile
 
 from octane.util import docker
 from octane.util import subprocess
@@ -63,3 +67,36 @@ def filter_members(archive, dir_name):
     for member in archive:
         if member.isfile() and member.name.startswith(dir_name):
             yield member
+
+
+@contextlib.contextmanager
+def update_cpio(img_path, dir_path=None):
+    tmp_dir = tempfile.mkdtemp(dir=dir_path)
+    try:
+        with subprocess.popen(
+                ["gunzip", "-c", img_path],
+                stdout=subprocess.PIPE) as proc:
+            subprocess.call(
+                ["cpio", "-id"], stdin=proc.stdout, cwd=tmp_dir)
+        yield tmp_dir
+        tmp_dir_len = len(tmp_dir)
+        with tempfile.NamedTemporaryFile(dir=dir_path) as new_img:
+            with subprocess.popen(
+                    ["cpio", "--format", "newc", "-o"],
+                    stdin=subprocess.PIPE,
+                    stdout=subprocess.PIPE,
+                    cwd=tmp_dir) as cpio:
+                with subprocess.popen(
+                        ["gzip", "-c"],
+                        stdin=cpio.stdout,
+                        stdout=new_img,
+                        cwd=tmp_dir):
+                    for path, dirs, files in os.walk(tmp_dir):
+                        for name in itertools.chain(dirs, files):
+                            p_name = os.path.join(path, name)[tmp_dir_len + 1:]
+                            cpio.stdin.write("{0}\n".format(p_name))
+                    cpio.stdin.close()
+            shutil.move(new_img.name, img_path)
+            new_img.delete = False
+    finally:
+        shutil.rmtree(tmp_dir)