Coverage for polar/benefit/strategies/discord/schemas.py: 94%

1from typing import Literal 

2 

3from pydantic import Field, computed_field, field_validator 

4 

5from polar.config import settings 

6from polar.kit import jwt 

7from polar.kit.schemas import Schema 

8from polar.models.benefit import BenefitType 

9 

10from ..base.schemas import ( 

11 BenefitBase, 

12 BenefitCreateBase, 

13 BenefitSubscriberBase, 

14 BenefitUpdateBase, 

15) 

16 

17 

18class BenefitDiscordProperties(Schema): 

19 """ 

20 Properties for a benefit of type `discord`. 

21 """ 

22 

23 guild_id: str = Field(..., description="The ID of the Discord server.") 

24 role_id: str = Field(..., description="The ID of the Discord role to grant.") 

25 kick_member: bool = Field( 

26 ..., 

27 description="Whether to kick the member from the Discord server on revocation.", 

28 ) 

29 

30 @computed_field # type: ignore[prop-decorator] 

31 @property 

32 def guild_token(self) -> str: 

33 return jwt.encode( 

34 data={"guild_id": self.guild_id}, 

35 secret=settings.SECRET, 

36 type="discord_guild_token", 

37 ) 

38 

39 

40class BenefitDiscordCreateProperties(Schema): 

41 """ 

42 Properties to create a benefit of type `discord`. 

43 """ 

44 

45 guild_token: str = Field(serialization_alias="guild_id") 

46 role_id: str = Field(..., description="The ID of the Discord role to grant.") 

47 kick_member: bool = Field( 

48 ..., 

49 description="Whether to kick the member from the Discord server on revocation.", 

50 ) 

51 

52 @field_validator("guild_token") 

53 @classmethod 

54 def validate_guild_token(cls, v: str) -> str: 

55 try: 

56 guild_token_data = jwt.decode( 

57 token=v, secret=settings.SECRET, type="discord_guild_token" 

58 ) 

59 return guild_token_data["guild_id"] 

60 except (KeyError, jwt.DecodeError, jwt.ExpiredSignatureError) as e: 

61 raise ValueError( 

62 "Invalid token. Please authenticate your Discord server again." 

63 ) from e 

64 

65 

66class BenefitDiscordSubscriberProperties(Schema): 

67 """ 

68 Properties available to subscribers for a benefit of type `discord`. 

69 """ 

70 

71 guild_id: str = Field(..., description="The ID of the Discord server.") 

72 

73 

74class BenefitDiscordCreate(BenefitCreateBase): 

75 type: Literal[BenefitType.discord] 

76 properties: BenefitDiscordCreateProperties 

77 

78 

79class BenefitDiscordUpdate(BenefitUpdateBase): 

80 type: Literal[BenefitType.discord] 

81 properties: BenefitDiscordCreateProperties | None = None 

82 

83 

84class BenefitDiscord(BenefitBase): 

85 """ 

86 A benefit of type `discord`. 

87 

88 Use it to automatically invite your backers to a Discord server. 

89 """ 

90 

91 type: Literal[BenefitType.discord] 

92 properties: BenefitDiscordProperties 

93 

94 

95class BenefitDiscordSubscriber(BenefitSubscriberBase): 

96 type: Literal[BenefitType.discord] 

97 properties: BenefitDiscordSubscriberProperties