Coverage for src/kwai/modules/identity/tokens/log_user_login_db_service.py: 100%

1"""Module that defines a LogUserLoginService with a database.""" 

2 

3from kwai.core.db.database import Database 

4from kwai.modules.identity.tokens.log_user_login_service import LogUserLoginService 

5from kwai.modules.identity.tokens.refresh_token import RefreshTokenEntity 

6from kwai.modules.identity.tokens.user_log import UserLogEntity 

7from kwai.modules.identity.tokens.user_log_db_repository import UserLogDbRepository 

8from kwai.modules.identity.tokens.value_objects import IpAddress, OpenId 

9from kwai.modules.identity.users.user_account import UserAccountEntity 

10 

11 

12class LogUserLoginDbService(LogUserLoginService): 

13 """Logs a login user request to the database.""" 

14 

15 def __init__( 

16 self, 

17 database: Database, 

18 *, 

19 email: str, 

20 client_ip: str, 

21 user_agent: str, 

22 open_id_sub: str = "", 

23 open_id_provider: str = "", 

24 ): 

25 self._db = database 

26 self._email = email 

27 if client_ip == "testclient": 

28 self._client_ip = IpAddress.create("127.0.0.1") 

29 else: 

30 self._client_ip = IpAddress.create(client_ip) 

31 self._user_agent = user_agent 

32 self._openId = OpenId(sub=open_id_sub, provider=open_id_provider) 

33 

34 async def notify_failure( 

35 self, 

36 message: str = "", 

37 *, 

38 user_account: UserAccountEntity | None = None, 

39 refresh_token: RefreshTokenEntity | None = None, 

40 ) -> None: 

41 repo = UserLogDbRepository(self._db) 

42 await repo.create( 

43 UserLogEntity( 

44 success=False, 

45 email=self._email, 

46 user_account=user_account, 

47 client_ip=self._client_ip, 

48 user_agent=self._user_agent, 

49 remark=message, 

50 openid=self._openId, 

51 ) 

52 ) 

53 

54 async def notify_success( 

55 self, 

56 *, 

57 user_account: UserAccountEntity | None = None, 

58 refresh_token: RefreshTokenEntity | None = None, 

59 ) -> None: 

60 repo = UserLogDbRepository(self._db) 

61 await repo.create( 

62 UserLogEntity( 

63 success=True, 

64 email=self._email, 

65 user_account=user_account, 

66 refresh_token=refresh_token, 

67 client_ip=self._client_ip, 

68 user_agent=self._user_agent, 

69 openid=self._openId, 

70 ) 

71 )