Coverage for jumpstarter_driver_flashers/bundle.py: 97%

1import os 

2from typing import Literal, Optional 

3 

4import yaml 

5from pydantic import BaseModel, Field 

6 

7 

8class FileAddress(BaseModel): 

9 file: str 

10 address: str 

11 

12 

13class DtbVariant(BaseModel): 

14 default: str 

15 address: str 

16 variants: dict[str, str] 

17 

18 

19class FlasherLogin(BaseModel): 

20 login_prompt: str 

21 username: str | None = None 

22 password: str | None = None 

23 prompt: str 

24 

25 

26class FlashBundleSpecV1Alpha1(BaseModel): 

27 manufacturer: str 

28 link: Optional[str] 

29 bootcmd: str 

30 shelltype: Literal["busybox"] = Field(default="busybox") 

31 login: FlasherLogin = Field(default_factory=lambda: FlasherLogin(login_prompt="login:", prompt="#")) 

32 default_target: str 

33 targets: dict[str, str] 

34 kernel: FileAddress 

35 initram: Optional[FileAddress] = None 

36 dtb: DtbVariant 

37 preflash_commands: list[str] = Field(default_factory=list) 

38 

39 

40class ObjectMeta(BaseModel): 

41 name: str 

42 

43 

44class FlasherBundleManifestV1Alpha1(BaseModel): 

45 apiVersion: Literal["jumpstarter.dev/v1alpha1"] = Field(default="jumpstarter.dev/v1alpha1") 

46 kind: Literal["FlashBundleManifest"] = Field(default="FlashBundleManifest") 

47 metadata: ObjectMeta 

48 spec: FlashBundleSpecV1Alpha1 

49 

50 def get_dtb_address(self) -> str: 

51 return self.spec.dtb.address 

52 

53 def get_dtb_file(self, variant: str | None = None) -> str | None: 

54 if not variant: 

55 variant = self.spec.dtb.default 

56 

57 return self.spec.dtb.variants.get(variant) 

58 

59 def get_kernel_address(self) -> str: 

60 return self.spec.kernel.address 

61 

62 def get_kernel_file(self) -> str: 

63 return self.spec.kernel.file 

64 

65 def get_initram_file(self) -> str | None: 

66 if not self.spec.initram: 

67 return None 

68 return self.spec.initram.file 

69 

70 def get_initram_address(self) -> str | None: 

71 if not self.spec.initram: 

72 return None 

73 return self.spec.initram.address 

74 

75 @classmethod 

76 def from_file(cls, path: os.PathLike): 

77 with open(path) as f: 

78 v = cls.model_validate(yaml.safe_load(f)) 

79 return v 

80 

81 @classmethod 

82 def from_string(cls, data: str): 

83 v = cls.model_validate(yaml.safe_load(data)) 

84 return v