Refactoring - most important change is to reuse method to create targ…

…et tar archive

docker_scripts/image.py

@@ -204,6 +204,10 @@ def _before_squashing(self):
         # Unpack exported image
         self._unpack(self.old_image_tar, self.old_image_dir)
 
+        # Remove the tar file early to save some space
+        self.log.debug("Removing exported tar (%s)..." % self.old_image_tar)
+        os.remove(self.old_image_tar)
+
         self.log.info("Squashing image '%s'..." % self.image)
 
     def _after_squashing(self):
@@ -215,11 +219,6 @@ def layer_paths(self):
         """
         pass
 
-    def unpack_image(self):
-        """
-        Unpacks old image.
-        """
-
     def export_tar_archive(self, target_tar_file):
         self._tar_image(target_tar_file, self.new_image_dir)
         self.log.info("Image available at '%s'" % target_tar_file)
@@ -234,6 +233,7 @@ def _files_in_layers(self, layers, directory):
         Prepare a list of files in all layers
         """
         files = {}
+
         for layer in layers:
             self.log.debug("Generating list of files in layer '%s'..." % layer)
             tar_file = os.path.join(directory, layer, "layer.tar")
@@ -259,19 +259,17 @@ def _prepare_tmp_directory(self, tmp_dir):
         return tmp_dir
 
     def _load_image(self, directory):
-        buf = six.BytesIO()
 
-        with tarfile.open(mode='w', fileobj=buf) as tar:
-            self.log.debug("Generating tar archive for the squashed image...")
-            with Chdir(directory):
-                tar.add(".")
-            self.log.debug("Archive generated")
+        tar_file = os.path.join(self.tmp_dir, "image.tar")
+
+        self._tar_image(tar_file, directory)
 
-        self.log.debug("Loading squashed image...")
-        self.docker.load_image(buf.getvalue())
-        self.log.debug("Image loaded!")
+        with open(tar_file, 'rb') as f:
+            self.log.debug("Loading squashed image...")
+            self.docker.load_image(f.read())
+            self.log.debug("Image loaded!")
 
-        buf.close()
+        os.remove(tar_file)
 
     def _tar_image(self, target_tar_file, directory):
         with tarfile.open(target_tar_file, 'w', format=tarfile.PAX_FORMAT) as tar:
@@ -340,10 +338,6 @@ def _unpack(self, tar_file, directory):
         with tarfile.open(tar_file, 'r') as tar:
             tar.extractall(path=directory)
 
-        # Remove the tar file early to save some space
-        self.log.debug("Removing exported tar (%s)..." % tar_file)
-        os.remove(tar_file)
-
         self.log.info("Archive unpacked!")
 
     def _read_layers(self, layers, image_id):
@@ -456,7 +450,7 @@ def _marker_files(self, tar, members):
 
         return marker_files
 
-    def _add_markers(self, markers, tar, layers_to_move, old_image_dir):
+    def _add_markers(self, markers, tar, files_in_layers):
         """
         This method is responsible for adding back all markers that were not
         added to the squashed layer AND files they refer to can be found in layers
@@ -470,9 +464,6 @@ def _add_markers(self, markers, tar, layers_to_move, old_image_dir):
             # No marker files to add
             return
 
-        # Find all files in layers that we don't squash
-        files_in_layers = self._files_in_layers(layers_to_move, old_image_dir)
-
         for marker, marker_file in six.iteritems(markers):
             actual_file = marker.name.replace('.wh.', '')
             should_be_added_back = False
@@ -568,7 +559,9 @@ def _squash_layers(self, layers_to_squash, layers_to_move):
                         # We added a file to the squashed tar, so let's note it
                         squashed_files.append(member.name)
 
-            self._add_markers(
-                missed_markers, squashed_tar, layers_to_move, self.old_image_dir)
+            # Find all files in layers that we don't squash
+            files_in_layers_to_move = self._files_in_layers(layers_to_move, self.old_image_dir)
+
+            self._add_markers(missed_markers, squashed_tar, files_in_layers_to_move)
 
         self.log.info("Squashing finished!")

tests/test_unit_v1_image.py

@@ -217,7 +217,7 @@ def setUp(self):
         self.squash = Image(self.log, self.docker_client, self.image, None)
 
     def test_should_not_fail_with_empty_list_of_markers_to_add(self):
-        self.squash._add_markers({}, None, None, None)
+        self.squash._add_markers({}, None, None)
 
     def test_should_add_all_marker_files_to_empty_tar(self):
         tar = mock.Mock()
@@ -227,7 +227,7 @@ def test_should_add_all_marker_files_to_empty_tar(self):
 
         markers = {marker_1: 'file'}
         with mock.patch('docker_scripts.image.Image._files_in_layers', return_value={}):
-            self.squash._add_markers(markers, tar, None, None)
+            self.squash._add_markers(markers, tar, None)
 
         self.assertTrue(len(tar.addfile.mock_calls) == 1)
         tar_info, marker_file = tar.addfile.call_args[0]
@@ -244,10 +244,9 @@ def test_should_skip_a_marker_file_if_file_is_in_unsquashed_layers(self):
         type(marker_2).name = mock.PropertyMock(return_value='.wh.marker_2')
 
         markers = {marker_1: 'file1', marker_2: 'file2'}
-        with mock.patch('docker_scripts.image.Image._files_in_layers', return_value={'1234layerdid': ['some/file', 'marker_1']}):
-            self.squash._add_markers(markers, tar, None, None)
+        self.squash._add_markers(markers, tar, {'1234layerdid': ['some/file', 'marker_1']})
 
-        self.assertTrue(len(tar.addfile.mock_calls) == 1)
+        self.assertEqual(len(tar.addfile.mock_calls), 1)
         tar_info, marker_file = tar.addfile.call_args[0]
         self.assertIsInstance(tar_info, tarfile.TarInfo)
         self.assertTrue(marker_file == 'file2')
@@ -262,8 +261,7 @@ def test_should_not_add_any_marker_files(self):
         type(marker_2).name = mock.PropertyMock(return_value='.wh.marker_2')
 
         markers = {marker_1: 'file1', marker_2: 'file2'}
-        with mock.patch('docker_scripts.image.Image._files_in_layers', return_value={'1234layerdid': ['some/file', 'marker_1', 'marker_2']}):
-            self.squash._add_markers(markers, tar, None, None)
+        self.squash._add_markers(markers, tar, {'1234layerdid': ['some/file', 'marker_1', 'marker_2']})
 
         self.assertTrue(len(tar.addfile.mock_calls) == 0)