Coverage for src/tests/fixtures/identity/tokens.py: 100%

1"""Module for defining factory fixtures for token entities.""" 

2 

3import pytest 

4 

5from kwai_bc_identity.tokens.access_token import AccessTokenEntity 

6from kwai_bc_identity.tokens.access_token_db_repository import ( 

7 AccessTokenDbRepository, 

8) 

9from kwai_bc_identity.tokens.refresh_token import RefreshTokenEntity 

10from kwai_bc_identity.tokens.refresh_token_db_repository import ( 

11 RefreshTokenDbRepository, 

12) 

13from kwai_bc_identity.tokens.token_identifier import TokenIdentifier 

14from kwai_bc_identity.users.user_account import UserAccountEntity 

15from kwai_core.db.uow import UnitOfWork 

16from kwai_core.domain.value_objects.timestamp import Timestamp 

17 

18 

19@pytest.fixture 

20def make_access_token(make_user_account): 

21 """A factory fixture for an access token.""" 

22 

23 def _make_access_token( 

24 identifier: TokenIdentifier | None = None, 

25 expiration: Timestamp | None = None, 

26 user_account: UserAccountEntity | None = None, 

27 ) -> AccessTokenEntity: 

28 return AccessTokenEntity( 

29 identifier=identifier or TokenIdentifier.generate(), 

30 expiration=expiration or Timestamp.create_with_delta(minutes=10), 

31 user_account=user_account or make_user_account(), 

32 ) 

33 

34 return _make_access_token 

35 

36 

37@pytest.fixture 

38async def make_access_token_in_db(database, make_access_token, make_user_account_in_db): 

39 """Factory fixture for creating an access token in database.""" 

40 created_entities = [] 

41 

42 async def _make_access_token_in_db( 

43 access_token: AccessTokenEntity | None = None, 

44 user_account: UserAccountEntity | None = None, 

45 ) -> AccessTokenEntity: 

46 if access_token is None: 

47 access_token = make_access_token( 

48 user_account=user_account or await make_user_account_in_db() 

49 ) 

50 repo = AccessTokenDbRepository(database) 

51 

52 async with UnitOfWork(database): 

53 access_token = await repo.create(access_token) 

54 

55 created_entities.append(access_token) 

56 

57 return access_token 

58 

59 yield _make_access_token_in_db 

60 

61 repo = AccessTokenDbRepository(database) 

62 for entity in created_entities: 

63 async with UnitOfWork(database): 

64 await repo.delete(entity) 

65 

66 

67@pytest.fixture 

68def make_refresh_token(make_access_token): 

69 """A factory fixture for refresh token.""" 

70 

71 def _make_refresh_token( 

72 access_token: AccessTokenEntity | None = None, 

73 identifier: TokenIdentifier | None = None, 

74 expiration: Timestamp | None = None, 

75 ) -> RefreshTokenEntity: 

76 return RefreshTokenEntity( 

77 access_token=access_token or make_access_token(), 

78 identifier=identifier or TokenIdentifier.generate(), 

79 expiration=expiration or Timestamp.create_with_delta(minutes=10), 

80 ) 

81 

82 return _make_refresh_token 

83 

84 

85@pytest.fixture 

86async def make_refresh_token_in_db( 

87 database, make_access_token_in_db, make_refresh_token 

88): 

89 """Factory fixture for creating a refresh token in database.""" 

90 created_entities = [] 

91 

92 async def _make_refresh_token_in_db( 

93 refresh_token: RefreshTokenEntity | None = None, 

94 ) -> RefreshTokenEntity: 

95 if refresh_token is None: 

96 refresh_token = make_refresh_token( 

97 access_token=await make_access_token_in_db() 

98 ) 

99 repo = RefreshTokenDbRepository(database) 

100 

101 async with UnitOfWork(database): 

102 access_token = await repo.create(refresh_token) 

103 

104 created_entities.append(access_token) 

105 

106 return access_token 

107 

108 yield _make_refresh_token_in_db 

109 

110 repo = AccessTokenDbRepository(database) 

111 for entity in created_entities: 

112 async with UnitOfWork(database): 

113 await repo.delete(entity)