Coverage for wrapper/mock/wrapper_mock.py: 86%

1from typing import Any 

2from typing import Dict 

3from typing import Optional 

4 

5from wrapper_model import ModelGeneration 

6 

7from PIL import Image 

8 

9 

10class MockGeneration(ModelGeneration): 

11 """Mock model generation for testing purposes.""" 

12 

13 def __init__(self) -> None: 

14 super().__init__("mock") 

15 

16 async def warmup(self) -> None: 

17 await self.generate() 

18 

19 async def generate( 

20 self, 

21 job_id: Optional[str] = None, 

22 *args: Any, 

23 **kwargs: Any 

24 ) -> Any: 

25 gen_timer = self._new_gen_timer(job_id) 

26 

27 self.running = True 

28 

29 try: 

30 if "output_type" not in kwargs: 

31 return None 

32 width = int(kwargs.get("width", 128)) 

33 height = int(kwargs.get("height", 64)) 

34 if kwargs["output_type"] == "pil": 

35 return Image.new("RGB", (width, height), color='blue') 

36 if kwargs["output_type"] == "jsonl": 

37 return '{"mock": "data"}' 

38 if kwargs["output_type"] == "audio_path": 

39 return "TODO audio_path" 

40 if kwargs["output_type"] == "video_binary": 

41 return "TODO video_binary" 

42 if kwargs["output_type"] == "video_path": 

43 return "TODO video_path" 

44 if kwargs["output_type"] == "tensor": 

45 return "TODO tensor" 

46 return kwargs["output_type"] 

47 finally: 

48 gen_timer.end("total") 

49 self.running = False 

50 

51 async def get_rest_args( 

52 self, 

53 data_json: Dict[str, str] 

54 ) -> Dict[str, Any]: 

55 return {}