1 files changed, 101 insertions, 1 deletions
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index 4715328592..e201b741c6 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -2334,7 +2334,7 @@ class TestFunctional(unittest.TestCase):
         image_fname = tools.GetOutputFilename('image.bin')
         image = Image.FromFile(image_fname)
         self.assertTrue(isinstance(image, Image))
-        self.assertEqual('image', image.image_name)
+        self.assertEqual('image', image.image_name[-5:])
 
     def testReadImageFail(self):
         """Test failing to read an image image's FDT map"""
@@ -2720,6 +2720,106 @@ class TestFunctional(unittest.TestCase):
         self.assertEqual(len(U_BOOT_DATA), entry.contents_size)
         self.assertEqual(len(U_BOOT_DATA), entry.size)
 
+    def _RunReplaceCmd(self, entry_name, data, decomp=True):
+        """Replace an entry in an image
+
+        This writes the entry data to update it, then opens the updated file and
+        returns the value that it now finds there.
+
+        Args:
+            entry_name: Entry name to replace
+            data: Data to replace it with
+            decomp: True to compress the data if needed, False if data is
+                already compressed so should be used as is
+
+        Returns:
+            Tuple:
+                data from entry
+                data from fdtmap (excluding header)
+        """
+        dtb_data = self._DoReadFileDtb('132_replace.dts', use_real_dtb=True,
+                                       update_dtb=True)[1]
+
+        self.assertIn('image', control.images)
+        image = control.images['image']
+        entries = image.GetEntries()
+        orig_dtb_data = entries['u-boot-dtb'].data
+        orig_fdtmap_data = entries['fdtmap'].data
+
+        image_fname = tools.GetOutputFilename('image.bin')
+        updated_fname = tools.GetOutputFilename('image-updated.bin')
+        tools.WriteFile(updated_fname, tools.ReadFile(image_fname))
+        control.WriteEntry(updated_fname, entry_name, data, decomp)
+        data = control.ReadEntry(updated_fname, entry_name, decomp)
+
+        # The DT data should not change
+        new_dtb_data = entries['u-boot-dtb'].data
+        self.assertEqual(new_dtb_data, orig_dtb_data)
+        new_fdtmap_data = entries['fdtmap'].data
+        self.assertEqual(new_fdtmap_data, orig_fdtmap_data)
+
+        return data, orig_fdtmap_data[fdtmap.FDTMAP_HDR_LEN:]
+
+    def testReplaceSimple(self):
+        """Test replacing a single file"""
+        expected = b'x' * len(U_BOOT_DATA)
+        data, expected_fdtmap = self._RunReplaceCmd('u-boot', expected)
+        self.assertEqual(expected, data)
+
+        # Test that the state looks right. There should be an FDT for the fdtmap
+        # that we jsut read back in, and it should match what we find in the
+        # 'control' tables. Checking for an FDT that does not exist should
+        # return None.
+        path, fdtmap = state.GetFdtContents('fdtmap')
+        self.assertIsNone(path)
+        self.assertEqual(expected_fdtmap, fdtmap)
+
+        dtb = state.GetFdtForEtype('fdtmap')
+        self.assertEqual(dtb.GetContents(), fdtmap)
+
+        missing_path, missing_fdtmap = state.GetFdtContents('missing')
+        self.assertIsNone(missing_path)
+        self.assertIsNone(missing_fdtmap)
+
+        missing_dtb = state.GetFdtForEtype('missing')
+        self.assertIsNone(missing_dtb)
+
+        self.assertEqual('/binman', state.fdt_path_prefix)
+
+    def testReplaceResizeFail(self):
+        """Test replacing a file by something larger"""
+        expected = U_BOOT_DATA + b'x'
+        with self.assertRaises(ValueError) as e:
+            self._RunReplaceCmd('u-boot', expected)
+        self.assertIn("Node '/u-boot': Entry data size does not match, but resize is disabled",
+                      str(e.exception))
+
+    def testReplaceMulti(self):
+        """Test replacing entry data where multiple images are generated"""
+        data = self._DoReadFileDtb('133_replace_multi.dts', use_real_dtb=True,
+                                   update_dtb=True)[0]
+        expected = b'x' * len(U_BOOT_DATA)
+        updated_fname = tools.GetOutputFilename('image-updated.bin')
+        tools.WriteFile(updated_fname, data)
+        entry_name = 'u-boot'
+        control.WriteEntry(updated_fname, entry_name, expected)
+        data = control.ReadEntry(updated_fname, entry_name)
+        self.assertEqual(expected, data)
+
+        # Check the state looks right.
+        self.assertEqual('/binman/image', state.fdt_path_prefix)
+
+        # Now check we can write the first image
+        image_fname = tools.GetOutputFilename('first-image.bin')
+        updated_fname = tools.GetOutputFilename('first-updated.bin')
+        tools.WriteFile(updated_fname, tools.ReadFile(image_fname))
+        entry_name = 'u-boot'
+        control.WriteEntry(updated_fname, entry_name, expected)
+        data = control.ReadEntry(updated_fname, entry_name)
+        self.assertEqual(expected, data)
+
+        # Check the state looks right.
+        self.assertEqual('/binman/first-image', state.fdt_path_prefix)
 
 if __name__ == "__main__":
     unittest.main()