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

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

2from abc import ABC, abstractmethod 

3from typing import AsyncIterator 

4 

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

6 

7 

8class CoachNotFoundException(Exception): 

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

10 

11 

12class CoachRepository(ABC): 

13 """Interface for a coach repository.""" 

14 

15 @abstractmethod 

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

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

18 

19 Args: 

20 id: The id of a coach. 

21 

22 Raises: 

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

24 exist. 

25 """ 

26 raise NotImplementedError 

27 

28 @abstractmethod 

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

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

31 

32 Args: 

33 id: A variable number of coach ids. 

34 """ 

35 raise NotImplementedError