Coverage for src/kwai/modules/training/trainings/training_repository.py: 100%

1"""Module that defines an interface for a repository of trainings.""" 

2 

3from abc import ABC, abstractmethod 

4from typing import AsyncIterator 

5 

6from kwai.modules.training.trainings.training import TrainingEntity, TrainingIdentifier 

7from kwai.modules.training.trainings.training_definition import TrainingDefinitionEntity 

8from kwai.modules.training.trainings.training_query import TrainingQuery 

9 

10 

11class TrainingNotFoundException(Exception): 

12 """Raised when a training can not be found.""" 

13 

14 

15class TrainingRepository(ABC): 

16 """An interface for a repository of trainings.""" 

17 

18 @abstractmethod 

19 def create_query(self) -> TrainingQuery: 

20 """Create a query for querying trainings.""" 

21 raise NotImplementedError 

22 

23 @abstractmethod 

24 def get_all( 

25 self, 

26 query: TrainingQuery | None = None, 

27 limit: int | None = None, 

28 offset: int | None = None, 

29 ) -> AsyncIterator[TrainingEntity]: 

30 """Return all trainings of a given query. 

31 

32 Args: 

33 query: The query to use for selecting the rows. 

34 limit: The maximum number of entities to return. 

35 offset: Skip the offset rows before beginning to return entities. 

36 

37 Returns: 

38 A list of trainings. 

39 """ 

40 raise NotImplementedError 

41 

42 @abstractmethod 

43 async def get_by_id(self, id: TrainingIdentifier) -> TrainingEntity: 

44 """Get the training with the given id. 

45 

46 Args: 

47 id: The id of the training. 

48 

49 Returns: 

50 A training entity. 

51 """ 

52 raise NotImplementedError 

53 

54 @abstractmethod 

55 async def create(self, training: TrainingEntity) -> TrainingEntity: 

56 """Save a training. 

57 

58 Args: 

59 training: The training to save. 

60 

61 Returns: 

62 A training entity with an identifier. 

63 """ 

64 raise NotImplementedError 

65 

66 @abstractmethod 

67 async def update(self, training: TrainingEntity) -> None: 

68 """Update a training. 

69 

70 Args: 

71 training: The training to save. 

72 """ 

73 raise NotImplementedError 

74 

75 @abstractmethod 

76 async def delete(self, training: TrainingEntity) -> None: 

77 """Delete a training. 

78 

79 Args: 

80 training: The training to delete. 

81 """ 

82 raise NotImplementedError 

83 

84 @abstractmethod 

85 async def reset_definition( 

86 self, training_definition: TrainingDefinitionEntity, delete: bool = False 

87 ) -> None: 

88 """Reset all trainings of a training definition. 

89 

90 Args: 

91 training_definition: The definition to use. 

92 delete: Delete training or update? 

93 

94 When delete is True, trainings will be deleted. When False 

95 the definition will be set to None. 

96 """ 

97 raise NotImplementedError