Coverage for bc/kwai-bc-training/src/kwai_bc_training/delete_training_schedule.py: 100%

1"""Module for the use case "Delete Training Schedule".""" 

2 

3from dataclasses import dataclass 

4 

5from kwai_bc_training.trainings.training_repository import TrainingRepository 

6from kwai_bc_training.trainings.training_schedule import ( 

7 TrainingScheduleIdentifier, 

8) 

9from kwai_bc_training.trainings.training_schedule_repository import ( 

10 TrainingScheduleRepository, 

11) 

12 

13 

14@dataclass(kw_only=True, frozen=True, slots=True) 

15class DeleteTrainingScheduleCommand: 

16 """Input for the use case.""" 

17 

18 id: int 

19 delete_trainings: bool = False 

20 

21 

22class DeleteTrainingSchedule: 

23 """Use case "Delete Training Schedule". 

24 

25 When delete_trainings is True, Trainings that are related to this schedule, 

26 will also be deleted. When False, the training will be updated. 

27 """ 

28 

29 def __init__( 

30 self, repo: TrainingScheduleRepository, training_repo: TrainingRepository 

31 ): 

32 """Initialize the use case. 

33 

34 Args: 

35 repo: A repository for deleting a training schedule. 

36 training_repo: A repository for deleting or updating a training(s). 

37 """ 

38 self._repo = repo 

39 self._training_repo = training_repo 

40 

41 async def execute(self, command: DeleteTrainingScheduleCommand): 

42 """Execute the use case. 

43 

44 Raises: 

45 TrainingSchemaNotFoundException: when the training schedule does not exist. 

46 """ 

47 training_schedule = await self._repo.get_by_id( 

48 TrainingScheduleIdentifier(command.id) 

49 ) 

50 

51 await self._training_repo.reset_schedule( 

52 training_schedule, command.delete_trainings 

53 ) 

54 

55 await self._repo.delete(training_schedule)