1010import  pytest 
1111
1212import  host_tools .drive  as  drive_tools 
13- from  framework .microvm  import  Microvm 
13+ from  framework .microvm  import  HugePagesConfig ,  Microvm 
1414
1515USEC_IN_MSEC  =  1000 
1616NS_IN_MSEC  =  1_000_000 
@@ -36,6 +36,7 @@ class SnapshotRestoreTest:
3636    nets : int  =  3 
3737    blocks : int  =  3 
3838    all_devices : bool  =  False 
39+     huge_pages : HugePagesConfig  =  HugePagesConfig .NONE 
3940
4041    @property  
4142    def  id (self ):
@@ -55,6 +56,7 @@ def configure_vm(self, microvm_factory, guest_kernel, rootfs, metrics) -> Microv
5556            vcpu_count = self .vcpus ,
5657            mem_size_mib = self .mem ,
5758            rootfs_io_engine = "Sync" ,
59+             huge_pages = self .huge_pages ,
5860        )
5961
6062        for  _  in  range (self .nets ):
@@ -134,11 +136,15 @@ def test_restore_latency(
134136
135137@pytest .mark .nonci  
136138@pytest .mark .parametrize ("uffd_handler" , [None , "valid" ]) 
139+ @pytest .mark .parametrize ("huge_pages" , HugePagesConfig ) 
137140def  test_post_restore_latency (
138-     microvm_factory , rootfs , guest_kernel_linux_5_10 , metrics , uffd_handler 
141+     microvm_factory , rootfs , guest_kernel_linux_5_10 , metrics , uffd_handler ,  huge_pages 
139142):
140143    """Collects latency metric of post-restore memory accesses done inside the guest""" 
141-     test_setup  =  SnapshotRestoreTest (mem = 1024 , vcpus = 2 )
144+     if  huge_pages  !=  HugePagesConfig .NONE  and  uffd_handler  is  None :
145+         pytest .skip ("huge page snapshots can only be restored using uffd" )
146+ 
147+     test_setup  =  SnapshotRestoreTest (mem = 1024 , vcpus = 2 , huge_pages = huge_pages )
142148    vm  =  test_setup .configure_vm (
143149        microvm_factory , guest_kernel_linux_5_10 , rootfs , metrics 
144150    )
0 commit comments