Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py: 31%

1import time 

2from collections.abc import Generator 

3 

4import click 

5 

6from .common import PowerReading 

7from jumpstarter.client import DriverClient 

8 

9 

10class PowerClient(DriverClient): 

11 def on(self) -> None: 

12 """Power on the device.""" 

13 self.call("on") 

14 

15 def off(self) -> None: 

16 """Power off the device.""" 

17 self.call("off") 

18 

19 def rescue(self) -> None: 

20 self.call("rescue") 

21 

22 def cycle(self, wait: int = 2): 

23 """Power cycle the device.""" 

24 self.logger.info("Starting power cycle sequence") 

25 self.off() 

26 self.logger.info(f"Waiting {wait} seconds...") 

27 time.sleep(wait) 

28 self.on() 

29 self.logger.info("Power cycle sequence complete") 

30 

31 def read(self) -> Generator[PowerReading, None, None]: 

32 """Read power data from the device.""" 

33 

34 for v in self.streamingcall("read"): 

35 yield PowerReading.model_validate(v, strict=True) 

36 

37 def cli(self): 

38 @click.group 

39 def base(): 

40 """Generic power""" 

41 pass 

42 

43 @base.command() 

44 def on(): 

45 """Power on""" 

46 self.on() 

47 

48 @base.command() 

49 def off(): 

50 """Power off""" 

51 self.off() 

52 

53 @base.command() 

54 @click.option("--wait", "-w", default=2, help="Wait time in seconds between off and on") 

55 def cycle(wait): 

56 """Power cycle""" 

57 click.echo(f"Power cycling with {wait} seconds wait time...") 

58 self.cycle(wait) 

59 

60 return base 

61 

62 

63class VirtualPowerClient(PowerClient): 

64 def off(self, destroy: bool = False) -> None: 

65 self.call('off', destroy) 

66 

67 def cli(self): 

68 parent = super().cli() 

69 

70 @parent.command(name='off') 

71 @click.option('--destroy', is_flag=True, help='destroy the instance after powering it off') 

72 def off(destroy: bool): 

73 """Power off""" 

74 self.off(destroy) 

75 

76 return parent