Coverage for jumpstarter_driver_flashers/bundle.py: 93%

1import os 

2from typing import Literal 

3 

4import yaml 

5from pydantic import BaseModel, Field 

6 

7 

8class FileAddress(BaseModel): 

9 file: str 

10 address: str 

11 

12class DtbVariant(BaseModel): 

13 bootcmd: None | str = None 

14 file: None | str = None 

15 

16class Dtb(BaseModel): 

17 default: str 

18 address: str 

19 variants: dict[str, DtbVariant] 

20 

21class FlasherLogin(BaseModel): 

22 login_prompt: str 

23 username: str | None = None 

24 password: str | None = None 

25 prompt: str 

26 

27 

28class FlashBundleSpecV1Alpha1(BaseModel): 

29 manufacturer: str 

30 link: None | str 

31 bootcmd: str 

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

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

34 default_target: str 

35 targets: dict[str, str] 

36 kernel: FileAddress 

37 initram: None | FileAddress = None 

38 dtb: None | Dtb = None 

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

40 

41 

42class ObjectMeta(BaseModel): 

43 name: str 

44 

45 

46class FlasherBundleManifestV1Alpha1(BaseModel): 

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

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

49 metadata: ObjectMeta 

50 spec: FlashBundleSpecV1Alpha1 

51 

52 def get_dtb_address(self) -> str | None: 

53 if not self.spec.dtb: 

54 return None 

55 return self.spec.dtb.address 

56 

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

58 if not self.spec.dtb: 

59 return None 

60 

61 # if no variant is provided, use the default variant name from the manifest 

62 if not variant: 

63 variant = self.spec.dtb.default 

64 

65 # look for the variant struct in this manifest 

66 variant_struct = self.spec.dtb.variants.get(variant) 

67 if variant_struct: 

68 return variant_struct.file 

69 else: 

70 raise ValueError(f"DTB variant {variant} not found in the manifest.") 

71 

72 def get_boot_cmd(self, variant: str | None = None) -> str: 

73 if not self.spec.dtb: 

74 return self.spec.bootcmd 

75 # if no variant is provided, use the default variant name from the manifest 

76 if not variant: 

77 variant = self.spec.dtb.default 

78 # look for the variant struct in this manifest 

79 variant_struct = self.spec.dtb.variants.get(variant) 

80 if variant_struct: 

81 # If variant has a custom bootcmd, use it; otherwise fall back to default 

82 if variant_struct.bootcmd: 

83 return variant_struct.bootcmd 

84 else: 

85 return self.spec.bootcmd 

86 else: 

87 raise ValueError(f"DTB variant {variant} not found in the manifest.") 

88 

89 def get_kernel_address(self) -> str: 

90 return self.spec.kernel.address 

91 

92 def get_kernel_file(self) -> str: 

93 return self.spec.kernel.file 

94 

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

96 if not self.spec.initram: 

97 return None 

98 return self.spec.initram.file 

99 

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

101 if not self.spec.initram: 

102 return None 

103 return self.spec.initram.address 

104 

105 @classmethod 

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

107 with open(path) as f: 

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

109 return v 

110 

111 @classmethod 

112 def from_string(cls, data: str): 

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

114 return v