Coverage for src/kwai/modules/training/coaches/coach_repository.py: 100%

1"""Module that defines an interface for a coach repository.""" 

2 

3from abc import ABC, abstractmethod 

4from typing import AsyncIterator 

5 

6from kwai.modules.training.coaches.coach import CoachEntity, CoachIdentifier 

7from kwai.modules.training.coaches.coach_query import CoachQuery 

8 

9 

10class CoachNotFoundException(Exception): 

11 """Raised when a coach is not found.""" 

12 

13 

14class CoachRepository(ABC): 

15 """Interface for a coach repository.""" 

16 

17 @abstractmethod 

18 def create_query(self) -> CoachQuery: 

19 """Create a coach query.""" 

20 raise NotImplementedError 

21 

22 @abstractmethod 

23 async def get_by_id(self, id: CoachIdentifier) -> CoachEntity: 

24 """Get the coach with the given id. 

25 

26 Args: 

27 id: The id of a coach. 

28 

29 Raises: 

30 CoachNotFoundException: raised when the coach with the given id does not 

31 exist. 

32 """ 

33 raise NotImplementedError 

34 

35 @abstractmethod 

36 async def get_by_ids(self, *id: CoachIdentifier) -> AsyncIterator[CoachEntity]: 

37 """Get all coaches for the given ids. 

38 

39 Args: 

40 id: A variable number of coach ids. 

41 """ 

42 raise NotImplementedError 

43 

44 @abstractmethod 

45 async def get_all( 

46 self, query: CoachQuery | None = None 

47 ) -> AsyncIterator[CoachEntity]: 

48 """Get all coaches. 

49 

50 Args: 

51 query: The query to use for getting all coaches. 

52 

53 When query is omitted, all coaches will be returned. 

54 """ 

55 raise NotImplementedError