BaseTools: FMMT: Add replacement regression testsHEADmaster

This change adds a few unit test cases to validate the FMMT tool.

Signed-off-by: Kun Qin <kun.qin@microsoft.com>

2 files changed, 291 insertions, 0 deletions

diff --git a/BaseTools/Tests/FMMTTests.py b/BaseTools/Tests/FMMTTests.py
new file mode 100644
index 0000000000..8905b4361e
--- /dev/null
+++ b/BaseTools/Tests/FMMTTests.py
@@ -0,0 +1,289 @@
+## @file

+# Unit tests for the Firmware Module Management Tool.

+#

+# Copyright (c) Microsoft Corporation.

+#

+# SPDX-License-Identifier: BSD-2-Clause-Patent

+##

+

+import logging

+import os

+import sys

+import tempfile

+import unittest

+import uuid

+from pathlib import Path

+from types import SimpleNamespace

+from unittest import mock

+

+import TestTools

+

+

+_TESTS_DIR = Path(__file__).resolve().parent

+_PYTHON_SOURCE = _TESTS_DIR.parent / "Source" / "Python"

+_FMMT_SOURCE = _PYTHON_SOURCE / "FMMT"

+sys.path.insert(0, str(_FMMT_SOURCE))

+sys.path.insert(0, str(_PYTHON_SOURCE))

+os.environ.setdefault("FmmtConfPath", "")

+

+# FMMT opens a build log when imported. Unit tests do not need a persistent log.

+with mock.patch("logging.FileHandler", return_value=logging.NullHandler()):

+    import core.FMMTOperation as FMMTOperation

+    from core.BiosTree import BIOSTREE

+    from core.BiosTree import FFS_FREE_SPACE

+    from core.BiosTree import FFS_TREE

+    from core.BiosTree import SEC_FV_TREE

+    from core.BiosTreeNode import FreeSpaceNode

+    from core.FvHandler import FvHandler

+    from FirmwareStorageFormat.Common import ModifyGuidFormat

+    from FirmwareStorageFormat.FfsFileHeader import EFI_FFS_FILE_HEADER

+

+

+class _FirmwareVolumeData:

+    def __init__(self, size, free_space=0, block_size=0x40):

+        self.Data = b""

+        self.Free_Space = free_space

+        self.Header = SimpleNamespace(

+            Attributes=0,

+            BlockMap=[SimpleNamespace(Length=block_size)],

+            ExtHeaderOffset=0,

+            FileSystemGuid=ModifyGuidFormat(

+                "00000000-0000-0000-0000-000000000000"

+            ),

+            FvLength=size,

+        )

+        self.Size = size

+

+    def ModCheckSum(self):

+        pass

+

+    def ModExtHeaderData(self):

+        pass

+

+    def ModFvExt(self):

+        pass

+

+    def ModFvSize(self):

+        pass

+

+

+class _TestFvHandler(FvHandler):

+    def CompressData(self, target_tree):

+        self.Status = True

+

+    def ModifyTest(self, parent_tree, needed_space):

+        self.Status = True

+

+

+class _TestRecursiveFvHandler(FvHandler):

+    def CompressData(self, target_tree):

+        self.Status = True

+

+

+def _create_ffs(name, size, data=b""):

+    node = BIOSTREE(name)

+    node.type = FFS_TREE

+    header = EFI_FFS_FILE_HEADER()

+    header.Name = ModifyGuidFormat(str(name))

+    header.Size[0] = size & 0xFF

+    header.Size[1] = (size >> 8) & 0xFF

+    header.Size[2] = (size >> 16) & 0xFF

+    node.Data = SimpleNamespace(

+        Data=data,

+        Header=header,

+        Name=name,

+        PadData=b"",

+        Size=size,

+    )

+    return node

+

+

+class FmmtReplaceTests(unittest.TestCase):

+    def test_replace_multiple_matches_uses_independent_ffs_trees(self):

+        target_guid = uuid.UUID("11111111-2222-3333-4444-555555555555")

+        targets = [SimpleNamespace(), SimpleNamespace()]

+        replacement = SimpleNamespace(

+            Data=SimpleNamespace(PadData=b"", Size=0x20),

+            Parent=None,

+        )

+

+        class _Root:

+            def __init__(self):

+                self.Findlist = []

+

+            def FindNode(self, name, matches):

+                matches.extend(targets)

+

+        class _Parser:

+            def __init__(self, root):

+                self.FinalData = b"output"

+                self.WholeFvTree = root

+

+            def Encapsulation(self, root, compress_status):

+                pass

+

+            def ParserFromRoot(self, root, data):

+                pass

+

+        input_parser = _Parser(_Root())

+        replacement_parser = _Parser(SimpleNamespace(Child=[replacement]))

+        replacements = []

+

+        class _Handler:

+            def __init__(self, new_ffs, target_ffs):

+                replacements.append(new_ffs)

+                new_ffs.Parent = target_ffs

+

+            def ReplaceFfs(self):

+                return True

+

+        with tempfile.TemporaryDirectory() as directory:

+            input_path = Path(directory) / "input.fd"

+            new_ffs_path = Path(directory) / "new.ffs"

+            output_path = Path(directory) / "output.fd"

+            input_path.write_bytes(b"input")

+            new_ffs_path.write_bytes(b"replacement")

+

+            with mock.patch.object(

+                FMMTOperation,

+                "FMMTParser",

+                side_effect=[input_parser, replacement_parser],

+            ), mock.patch.object(FMMTOperation, "FvHandler", _Handler):

+                FMMTOperation.ReplaceFfs(

+                    str(input_path),

+                    target_guid,

+                    str(new_ffs_path),

+                    str(output_path),

+                )

+

+        self.assertEqual(2, len(replacements))

+        self.assertIsNot(replacements[0], replacements[1])

+        self.assertIs(targets[0], replacements[0].Parent)

+        self.assertIs(targets[1], replacements[1].Parent)

+

+    def test_smaller_replacement_creates_free_space_node(self):

+        prefix = _create_ffs(

+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18

+        )

+        target = _create_ffs(

+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x40

+        )

+        replacement = _create_ffs(target.Data.Name, 0x18)

+        volume = BIOSTREE("FV")

+        volume.type = SEC_FV_TREE

+        volume.Data = _FirmwareVolumeData(0x80)

+        volume.insertChild(prefix)

+        volume.insertChild(target)

+

+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

+

+        self.assertIs(replacement, volume.Child[1])

+        free_space = volume.Child[-1]

+        self.assertEqual(FFS_FREE_SPACE, free_space.type)

+        self.assertIsInstance(free_space.Data, FreeSpaceNode)

+        self.assertEqual(b"\xFF" * 0x28, free_space.Data.Data)

+

+    def test_growing_nested_fv_preserves_final_ffs(self):

+        prefix = _create_ffs(

+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"

+        )

+        target = _create_ffs(

+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"

+        )

+        trailing = _create_ffs(

+            uuid.UUID("30000000-0000-0000-0000-000000000003"),

+            0x18,

+            b"trailing",

+        )

+        replacement = _create_ffs(target.Data.Name, 0x30, b"replacement")

+        volume = BIOSTREE("FV")

+        volume.type = SEC_FV_TREE

+        volume.Data = _FirmwareVolumeData(0x80)

+        volume.insertChild(prefix)

+        volume.insertChild(target)

+        volume.insertChild(trailing)

+

+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

+

+        self.assertEqual(b"trailing", trailing.Data.Data)

+        self.assertIn(trailing, volume.Child)

+        self.assertIs(replacement, volume.Child[1])

+        free_space = volume.Child[-1]

+        self.assertEqual(FFS_FREE_SPACE, free_space.type)

+        self.assertIsInstance(free_space.Data, FreeSpaceNode)

+        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)

+

+    def test_block_aligned_growth_of_only_ffs_does_not_crash(self):

+        target = _create_ffs(

+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"

+        )

+        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")

+        volume = BIOSTREE("FV")

+        volume.type = SEC_FV_TREE

+        volume.Data = _FirmwareVolumeData(0x40)

+        volume.insertChild(target)

+

+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

+

+        self.assertEqual([replacement], volume.Child)

+        self.assertEqual(0, volume.Data.Free_Space)

+

+    def test_block_aligned_growth_preserves_ffs_order(self):

+        prefix = _create_ffs(

+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"

+        )

+        target = _create_ffs(

+            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"

+        )

+        trailing = _create_ffs(

+            uuid.UUID("30000000-0000-0000-0000-000000000003"),

+            0x18,

+            b"trailing",

+        )

+        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")

+        volume = BIOSTREE("FV")

+        volume.type = SEC_FV_TREE

+        volume.Data = _FirmwareVolumeData(0x80)

+        volume.insertChild(prefix)

+        volume.insertChild(target)

+        volume.insertChild(trailing)

+

+        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

+

+        self.assertEqual([prefix, replacement, trailing], volume.Child)

+        self.assertEqual(0, volume.Data.Free_Space)

+

+    def test_growing_ancestor_fv_preserves_final_ffs(self):

+        prefix = _create_ffs(

+            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"

+        )

+        trailing = _create_ffs(

+            uuid.UUID("30000000-0000-0000-0000-000000000003"),

+            0x18,

+            b"trailing",

+        )

+        volume = BIOSTREE("FV")

+        volume.type = SEC_FV_TREE

+        volume.Data = _FirmwareVolumeData(0x80)

+        volume.insertChild(prefix)

+        volume.insertChild(trailing)

+

+        parent = BIOSTREE("ROOT")

+        parent.type = FMMTOperation.ROOT_TREE

+        parent.insertChild(volume)

+

+        _TestRecursiveFvHandler(None).ModifyTest(volume, 0x10)

+

+        self.assertEqual(b"trailing", trailing.Data.Data)

+        self.assertIn(trailing, volume.Child)

+        free_space = volume.Child[-1]

+        self.assertEqual(FFS_FREE_SPACE, free_space.type)

+        self.assertIsInstance(free_space.Data, FreeSpaceNode)

+        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)

+

+

+TheTestSuite = TestTools.MakeTheTestSuite(locals())

+

+

+if __name__ == "__main__":

+    unittest.TextTestRunner(verbosity=2).run(TheTestSuite())

diff --git a/BaseTools/Tests/PythonToolsTests.py b/BaseTools/Tests/PythonToolsTests.py
index 95b3202a48..561bc367f2 100644
--- a/BaseTools/Tests/PythonToolsTests.py
+++ b/BaseTools/Tests/PythonToolsTests.py
@@ -22,6 +22,8 @@ def TheTestSuite():
     suites.append(CheckUnicodeSourceFiles.TheTestSuite())

     import MetaFileParserTests

     suites.append(MetaFileParserTests.TheTestSuite())

+    import FMMTTests

+    suites.append(FMMTTests.TheTestSuite())

     return unittest.TestSuite(suites)

 

 if __name__ == '__main__':