refactor(test): introduce utils_uffd.py

Currently, all tests that want to make use of UFFD import UFFD related
functions from test_uffd.py. That's a bit awkward, so move them into
their own utility module.

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

tests/framework/utils.py

@@ -10,7 +10,6 @@
 import re
 import select
 import signal
-import stat
 import subprocess
 import time
 import typing
@@ -133,69 +132,6 @@ def chroot(path):
         os.chdir(working_dir)
 
 
-class UffdHandler:
-    """Describe the UFFD page fault handler process."""
-
-    def __init__(self, name, socket_path, mem_file, chroot_path, log_file_name):
-        """Instantiate the handler process with arguments."""
-        self._proc = None
-        self._handler_name = name
-        self._socket_path = socket_path
-        self._mem_file = mem_file
-        self._chroot = chroot_path
-        self._log_file = log_file_name
-
-    def spawn(self, uid, gid):
-        """Spawn handler process using arguments provided."""
-
-        with chroot(self._chroot):
-            st = os.stat(self._handler_name)
-            os.chmod(self._handler_name, st.st_mode | stat.S_IEXEC)
-
-            chroot_log_file = Path("/") / self._log_file
-            with open(chroot_log_file, "w", encoding="utf-8") as logfile:
-                args = [f"/{self._handler_name}", self._socket_path, self._mem_file]
-                self._proc = subprocess.Popen(
-                    args, stdout=logfile, stderr=subprocess.STDOUT
-                )
-
-            # Give it time start and fail, if it really has too (bad things happen).
-            time.sleep(1)
-            if not self.is_running():
-                print(chroot_log_file.read_text(encoding="utf-8"))
-                assert False, "Could not start PF handler!"
-
-            # The page fault handler will create the socket path with root rights.
-            # Change rights to the jailer's.
-            os.chown(self._socket_path, uid, gid)
-
-    @property
-    def proc(self):
-        """Return UFFD handler process."""
-        return self._proc
-
-    def is_running(self):
-        """Check if UFFD process is running"""
-        return self.proc is not None and self.proc.poll() is None
-
-    @property
-    def log_file(self):
-        """Return the path to the UFFD handler's log file"""
-        return Path(self._chroot) / Path(self._log_file)
-
-    @property
-    def log_data(self):
-        """Return the log data of the UFFD handler"""
-        if self.log_file is None:
-            return ""
-        return self.log_file.read_text(encoding="utf-8")
-
-    def __del__(self):
-        """Tear down the UFFD handler process."""
-        if self.proc is not None:
-            self.proc.kill()
-
-
 class CpuMap:
     """Cpu map from real cpu cores to containers visible cores.
 

tests/framework/utils_uffd.py

@@ -0,0 +1,92 @@
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""UFFD related utility functions"""
+
+import os
+import stat
+import subprocess
+import time
+from pathlib import Path
+
+from framework.utils import chroot
+
+SOCKET_PATH = "/firecracker-uffd.sock"
+
+
+class UffdHandler:
+    """Describe the UFFD page fault handler process."""
+
+    def __init__(self, name, socket_path, mem_file, chroot_path, log_file_name):
+        """Instantiate the handler process with arguments."""
+        self._proc = None
+        self._handler_name = name
+        self._socket_path = socket_path
+        self._mem_file = mem_file
+        self._chroot = chroot_path
+        self._log_file = log_file_name
+
+    def spawn(self, uid, gid):
+        """Spawn handler process using arguments provided."""
+
+        with chroot(self._chroot):
+            st = os.stat(self._handler_name)
+            os.chmod(self._handler_name, st.st_mode | stat.S_IEXEC)
+
+            chroot_log_file = Path("/") / self._log_file
+            with open(chroot_log_file, "w", encoding="utf-8") as logfile:
+                args = [f"/{self._handler_name}", self._socket_path, self._mem_file]
+                self._proc = subprocess.Popen(
+                    args, stdout=logfile, stderr=subprocess.STDOUT
+                )
+
+            # Give it time start and fail, if it really has too (bad things happen).
+            time.sleep(1)
+            if not self.is_running():
+                print(chroot_log_file.read_text(encoding="utf-8"))
+                assert False, "Could not start PF handler!"
+
+            # The page fault handler will create the socket path with root rights.
+            # Change rights to the jailer's.
+            os.chown(self._socket_path, uid, gid)
+
+    @property
+    def proc(self):
+        """Return UFFD handler process."""
+        return self._proc
+
+    def is_running(self):
+        """Check if UFFD process is running"""
+        return self.proc is not None and self.proc.poll() is None
+
+    @property
+    def log_file(self):
+        """Return the path to the UFFD handler's log file"""
+        return Path(self._chroot) / Path(self._log_file)
+
+    @property
+    def log_data(self):
+        """Return the log data of the UFFD handler"""
+        if self.log_file is None:
+            return ""
+        return self.log_file.read_text(encoding="utf-8")
+
+    def __del__(self):
+        """Tear down the UFFD handler process."""
+        if self.proc is not None:
+            self.proc.kill()
+
+
+def spawn_pf_handler(vm, handler_path, mem_path):
+    """Spawn page fault handler process."""
+    # Copy snapshot memory file into chroot of microVM.
+    jailed_mem = vm.create_jailed_resource(mem_path)
+    # Copy the valid page fault binary into chroot of microVM.
+    jailed_handler = vm.create_jailed_resource(handler_path)
+    handler_name = os.path.basename(jailed_handler)
+
+    uffd_handler = UffdHandler(
+        handler_name, SOCKET_PATH, jailed_mem, vm.chroot(), "uffd.log"
+    )
+    uffd_handler.spawn(vm.jailer.uid, vm.jailer.gid)
+
+    return uffd_handler

tests/integration_tests/functional/test_uffd.py

@@ -8,9 +8,8 @@
 import pytest
 import requests
 
-from framework.utils import Timeout, UffdHandler, check_output
-
-SOCKET_PATH = "/firecracker-uffd.sock"
+from framework.utils import Timeout, check_output
+from framework.utils_uffd import SOCKET_PATH, spawn_pf_handler
 
 
 @pytest.fixture(scope="function", name="snapshot")
@@ -36,22 +35,6 @@ def snapshot_fxt(microvm_factory, guest_kernel_linux_5_10, rootfs):
     yield snapshot
 
 
-def spawn_pf_handler(vm, handler_path, mem_path):
-    """Spawn page fault handler process."""
-    # Copy snapshot memory file into chroot of microVM.
-    jailed_mem = vm.create_jailed_resource(mem_path)
-    # Copy the valid page fault binary into chroot of microVM.
-    jailed_handler = vm.create_jailed_resource(handler_path)
-    handler_name = os.path.basename(jailed_handler)
-
-    uffd_handler = UffdHandler(
-        handler_name, SOCKET_PATH, jailed_mem, vm.chroot(), "uffd.log"
-    )
-    uffd_handler.spawn(vm.jailer.uid, vm.jailer.gid)
-
-    return uffd_handler
-
-
 def test_bad_socket_path(uvm_plain, snapshot):
     """
     Test error scenario when socket path does not exist.

tests/integration_tests/performance/test_huge_pages.py

@@ -10,7 +10,7 @@
 from framework.microvm import HugePagesConfig
 from framework.properties import global_props
 from framework.utils_ftrace import ftrace_events
-from integration_tests.functional.test_uffd import SOCKET_PATH, spawn_pf_handler
+from framework.utils_uffd import SOCKET_PATH, spawn_pf_handler
 
 
 def check_hugetlbfs_in_use(pid: int, allocation_name: str):