tests: add helper method MicrovmFactory.build_from_snapshot

This is a fairly common repeated pattern, so extract it into a method to
make writing tests simpler.

Signed-off-by: Pablo Barbáchano <[email protected]>

Author: pb8o

tests/conftest.py

@@ -477,9 +477,7 @@ def uvm_restored(microvm_factory, guest_kernel, rootfs, cpu_template):
     uvm = uvm_booted(microvm_factory, guest_kernel, rootfs, cpu_template)
     snapshot = uvm.snapshot_full()
     uvm.kill()
-    uvm2 = microvm_factory.build()
-    uvm2.spawn()
-    uvm2.restore_from_snapshot(snapshot, resume=True)
+    uvm2 = microvm_factory.build_from_snapshot(snapshot)
     uvm2.cpu_template_name = uvm.cpu_template_name
     return uvm2
 

tests/framework/microvm.py

@@ -1084,6 +1084,13 @@ def build(self, kernel=None, rootfs=None, **kwargs):
             vm.ssh_key = ssh_key
         return vm
 
+    def build_from_snapshot(self, snapshot: Snapshot):
+        """Build a microvm from a snapshot"""
+        vm = self.build()
+        vm.spawn()
+        vm.restore_from_snapshot(snapshot, resume=True)
+        return vm
+
     def kill(self):
         """Clean up all built VMs"""
         for vm in self.vms:

tests/integration_tests/functional/test_api.py

@@ -1166,10 +1166,7 @@ def test_get_full_config_after_restoring_snapshot(microvm_factory, uvm_nano):
     ]
 
     snapshot = uvm_nano.snapshot_full()
-    uvm2 = microvm_factory.build()
-    uvm2.spawn()
-    uvm2.restore_from_snapshot(snapshot, resume=True)
-
+    uvm2 = microvm_factory.build_from_snapshot(snapshot)
     expected_cfg = setup_cfg.copy()
 
     # We expect boot-source to be set with the following values

tests/integration_tests/functional/test_balloon.py

@@ -484,9 +484,7 @@ def test_balloon_snapshot(microvm_factory, guest_kernel, rootfs):
     assert first_reading > second_reading
 
     snapshot = vm.snapshot_full()
-    microvm = microvm_factory.build()
-    microvm.spawn()
-    microvm.restore_from_snapshot(snapshot, resume=True)
+    microvm = microvm_factory.build_from_snapshot(snapshot)
 
     # Get the firecracker from snapshot pid, and open an ssh connection.
     firecracker_pid = microvm.firecracker_pid

tests/integration_tests/functional/test_cmd_line_parameters.py

@@ -96,9 +96,7 @@ def test_cli_metrics_if_resume_no_metrics(uvm_plain, microvm_factory):
     snapshot = uvm1.snapshot_full()
 
     # When: restoring from the snapshot
-    uvm2 = microvm_factory.build()
-    uvm2.spawn()
-    uvm2.restore_from_snapshot(snapshot)
+    uvm2 = microvm_factory.build_from_snapshot(snapshot)
 
     # Then: the old metrics configuration does not exist
     metrics2 = Path(uvm2.jailer.chroot_path()) / metrics_path.name

tests/integration_tests/functional/test_snapshot_basic.py

@@ -50,39 +50,25 @@ def _get_guest_drive_size(ssh_connection, guest_dev_name="/dev/vdb"):
     return lines[1].strip()
 
 
-def test_resume_after_restoration(uvm_nano, microvm_factory):
-    """Tests snapshot is resumable after restoration.
+@pytest.mark.parametrize("resume_at_restore", [True, False])
+def test_resume(uvm_nano, microvm_factory, resume_at_restore):
+    """Tests snapshot is resumable at or after restoration.
 
-    Check that a restored microVM is resumable by calling PATCH /vm with Resumed
-    after PUT /snapshot/load with `resume_vm=False`.
+    Check that a restored microVM is resumable by either
+    a. PUT /snapshot/load with `resume_vm=False`, then calling PATCH /vm resume=True
+    b. PUT /snapshot/load with `resume_vm=True`
     """
     vm = uvm_nano
     vm.add_net_iface()
     vm.start()
-
-    snapshot = vm.snapshot_full()
-
-    restored_vm = microvm_factory.build()
-    restored_vm.spawn()
-    restored_vm.restore_from_snapshot(snapshot)
-    restored_vm.resume()
-
-
-def test_resume_at_restoration(uvm_nano, microvm_factory):
-    """Tests snapshot is resumable at restoration.
-
-    Check that a restored microVM is resumable by calling PUT /snapshot/load
-    with `resume_vm=True`.
-    """
-    vm = uvm_nano
-    vm.add_net_iface()
-    vm.start()
-
     snapshot = vm.snapshot_full()
-
     restored_vm = microvm_factory.build()
     restored_vm.spawn()
-    restored_vm.restore_from_snapshot(snapshot, resume=True)
+    restored_vm.restore_from_snapshot(snapshot, resume=resume_at_restore)
+    if not resume_at_restore:
+        assert restored_vm.state == "Paused"
+        restored_vm.resume()
+    assert restored_vm.state == "Running"
 
 
 def test_snapshot_current_version(uvm_nano):
@@ -228,9 +214,7 @@ def test_patch_drive_snapshot(uvm_nano, microvm_factory):
 
     # Load snapshot in a new Firecracker microVM.
     logger.info("Load snapshot, mem %s", snapshot.mem)
-    vm = microvm_factory.build()
-    vm.spawn()
-    vm.restore_from_snapshot(snapshot, resume=True)
+    vm = microvm_factory.build_from_snapshot(snapshot)
 
     # Attempt to connect to resumed microvm and verify the new microVM has the
     # right scratch drive.
@@ -319,9 +303,7 @@ def test_negative_postload_api(uvm_plain, microvm_factory):
     basevm.kill()
 
     # Do not resume, just load, so we can still call APIs that work.
-    microvm = microvm_factory.build()
-    microvm.spawn()
-    microvm.restore_from_snapshot(snapshot, resume=True)
+    microvm = microvm_factory.build_from_snapshot(snapshot)
 
     fail_msg = "The requested operation is not supported after starting the microVM"
     with pytest.raises(RuntimeError, match=fail_msg):
@@ -486,9 +468,7 @@ def test_diff_snapshot_overlay(guest_kernel, rootfs, microvm_factory):
 
     assert not filecmp.cmp(merged_snapshot.mem, first_snapshot_backup, shallow=False)
 
-    new_vm = microvm_factory.build()
-    new_vm.spawn()
-    new_vm.restore_from_snapshot(merged_snapshot, resume=True)
+    _ = microvm_factory.build_from_snapshot(merged_snapshot)
 
     # Check that the restored VM works
 
@@ -510,9 +490,7 @@ def test_snapshot_overwrite_self(guest_kernel, rootfs, microvm_factory):
     snapshot = base_vm.snapshot_full()
     base_vm.kill()
 
-    vm = microvm_factory.build()
-    vm.spawn()
-    vm.restore_from_snapshot(snapshot, resume=True)
+    vm = microvm_factory.build_from_snapshot(snapshot)
 
     # When restoring a snapshot, vm.restore_from_snapshot first copies
     # the memory file (inside of the jailer) to /mem.src

tests/integration_tests/functional/test_snapshot_editor.py

@@ -68,6 +68,4 @@ def test_remove_regs(uvm_nano, microvm_factory):
     assert MIDR_EL1 not in stdout
 
     # test that we can restore from a snapshot
-    new_vm = microvm_factory.build()
-    new_vm.spawn()
-    new_vm.restore_from_snapshot(snapshot, resume=True)
+    _ = microvm_factory.build_from_snapshot(snapshot)

tests/integration_tests/functional/test_snapshot_not_losing_dirty_pages.py

@@ -63,9 +63,6 @@ def test_diff_snapshot_works_after_error(
 
     # Now there is enough space for it to work
     snap2 = uvm.snapshot_diff()
-
-    vm2 = microvm_factory.build()
-    vm2.spawn()
-    vm2.restore_from_snapshot(snap2, resume=True)
-
     uvm.kill()
+
+    _vm2 = microvm_factory.build_from_snapshot(snap2)

tests/integration_tests/functional/test_vsock.py

@@ -199,10 +199,7 @@ def test_vsock_transport_reset_h2g(
     test_vm.kill()
 
     # Load snapshot.
-
-    vm2 = microvm_factory.build()
-    vm2.spawn()
-    vm2.restore_from_snapshot(snapshot, resume=True)
+    vm2 = microvm_factory.build_from_snapshot(snapshot)
 
     # Check that vsock device still works.
     # Test guest-initiated connections.
@@ -231,9 +228,7 @@ def test_vsock_transport_reset_g2h(uvm_nano, microvm_factory):
 
     for _ in range(5):
         # Load snapshot.
-        new_vm = microvm_factory.build()
-        new_vm.spawn()
-        new_vm.restore_from_snapshot(snapshot, resume=True)
+        new_vm = microvm_factory.build_from_snapshot(snapshot)
 
         # After snap restore all vsock connections should be
         # dropped. This means guest socat should exit same way