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

1"""Module for defining an abstract class for importing member entities.""" 

2 

3from abc import ABC, abstractmethod 

4from dataclasses import dataclass 

5from typing import AsyncGenerator 

6 

7from async_lru import alru_cache 

8 

9from kwai.core.domain.value_objects.owner import Owner 

10from kwai.modules.club.domain.country import CountryEntity 

11from kwai.modules.club.domain.file_upload import FileUploadEntity 

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

13from kwai.modules.club.repositories.country_repository import CountryRepository 

14 

15 

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

17class Result: 

18 """Base dataclass for a result of a member import.""" 

19 

20 row: int 

21 

22 

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

24class OkResult(Result): 

25 """Dataclass for a successful member import.""" 

26 

27 member: MemberEntity 

28 

29 

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

31class FailureResult(Result): 

32 """Dataclass for a failed member import.""" 

33 

34 message: str 

35 

36 

37class MemberImporter(ABC): 

38 """Abstract class for importing member entities from a file.""" 

39 

40 def __init__(self, filename: str, owner: Owner, country_repo: CountryRepository): 

41 """Initialize the importer. 

42 

43 Args: 

44 filename: The name of the csv file. 

45 owner: The user that started the upload. 

46 country_repo: A repository to get the nationality of a member. 

47 """ 

48 self._filename = filename 

49 self._owner = owner 

50 self._country_repo = country_repo 

51 

52 @abstractmethod 

53 def import_(self) -> AsyncGenerator[Result, None]: 

54 """Import member entities. 

55 

56 For each imported (or failed import) of a member, a result will be yielded. 

57 """ 

58 

59 def create_file_upload_entity(self, preview: bool) -> FileUploadEntity: 

60 """Create a file upload entity.""" 

61 return FileUploadEntity( 

62 filename=self._filename, owner=self._owner, preview=preview 

63 ) 

64 

65 @staticmethod 

66 @alru_cache 

67 async def _get_country( 

68 country_repo: CountryRepository, iso_2: str 

69 ) -> CountryEntity: 

70 """Gets the country from the repository. 

71 

72 The value is cached. When the country does not exist, a CountryNotFoundException 

73 is raised. 

74 

75 Note: to avoid memory leaks, this method is a static method. 

76 """ 

77 return await country_repo.get_by_iso_2(iso_2)