Skip to content

Commit d76e25f

Browse files
committed
refactor(test): Stop SnapshotType being a str enum
With support for "mincore diff snapshots", we'll have two snapshot types that map to "diff" in the firecracker API, so the enum can no longer be string based. Instead just have a property that translates to Firecracker API types. Signed-off-by: Patrick Roy <roypat@amazon.co.uk>
1 parent d7f6f88 commit d76e25f

File tree

1 file changed

+13
-4
lines changed

1 file changed

+13
-4
lines changed

tests/framework/microvm.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import uuid
2323
from collections import namedtuple
2424
from dataclasses import dataclass
25-
from enum import Enum
25+
from enum import Enum, auto
2626
from functools import lru_cache
2727
from pathlib import Path
2828
from typing import Optional
@@ -49,8 +49,8 @@
4949
class SnapshotType(Enum):
5050
"""Supported snapshot types."""
5151

52-
FULL = "Full"
53-
DIFF = "Diff"
52+
FULL = auto()
53+
DIFF = auto()
5454

5555
def __repr__(self):
5656
cls_name = self.__class__.__name__
@@ -66,6 +66,15 @@ def needs_dirty_page_tracking(self) -> bool:
6666
"""Does taking this snapshot type require dirty page tracking to be enabled?"""
6767
return self == SnapshotType.DIFF
6868

69+
@property
70+
def api_type(self) -> str:
71+
"""Converts this `SnapshotType` to the string value expected by the Firecracker API"""
72+
match self:
73+
case SnapshotType.FULL:
74+
return "Full"
75+
case SnapshotType.DIFF:
76+
return "Diff"
77+
6978

7079
def hardlink_or_copy(src, dst):
7180
"""If src and dst are in the same device, hardlink. Otherwise, copy."""
@@ -984,7 +993,7 @@ def make_snapshot(
984993
self.api.snapshot_create.put(
985994
mem_file_path=str(mem_path),
986995
snapshot_path=str(vmstate_path),
987-
snapshot_type=snapshot_type.value,
996+
snapshot_type=snapshot_type.api_type,
988997
)
989998
root = Path(self.chroot())
990999
return Snapshot(

0 commit comments

Comments
 (0)