Coverage for src/kwai/modules/club/repositories/file_upload_db_repository.py: 100%

1"""Module that implements a FileUploadRepository using a database.""" 

2 

3from kwai.core.db.database import Database 

4from kwai.core.domain.entity import Entity 

5from kwai.modules.club.domain.file_upload import FileUploadEntity, FileUploadIdentifier 

6from kwai.modules.club.domain.member import MemberEntity 

7from kwai.modules.club.repositories._tables import FileUploadRow, MemberUploadRow 

8from kwai.modules.club.repositories.file_upload_repository import ( 

9 DuplicateMemberUploadedException, 

10 FileUploadRepository, 

11) 

12 

13 

14class FileUploadDbRepository(FileUploadRepository): 

15 """An implementation of a FileUploadRepository using a database.""" 

16 

17 def __init__(self, database: Database): 

18 """Initialize the repository. 

19 

20 Args: 

21 database: The database for this repository. 

22 """ 

23 self._database = database 

24 self._saved_members: dict[str, MemberEntity] = {} 

25 

26 async def create(self, file_upload: FileUploadEntity) -> FileUploadEntity: 

27 new_id = await self._database.insert( 

28 FileUploadRow.__table_name__, FileUploadRow.persist(file_upload) 

29 ) 

30 return Entity.replace(file_upload, id_=FileUploadIdentifier(new_id)) 

31 

32 def is_duplicate(self, member: MemberEntity) -> bool: 

33 return member.license.number in self._saved_members 

34 

35 async def save_member(self, file_upload: FileUploadEntity, member: MemberEntity): 

36 if self.is_duplicate(member): 

37 raise DuplicateMemberUploadedException( 

38 f"Member with license {member.license.number} is already uploaded." 

39 ) 

40 self._saved_members[member.license.number] = member 

41 await self._database.insert( 

42 MemberUploadRow.__table_name__, 

43 MemberUploadRow.persist(file_upload, member), 

44 )