Coverage for src/kwai/api/v1/trainings/schemas/training_definition.py: 93%

1"""Module for the training definition schema.""" 

2 

3from pydantic import BaseModel 

4 

5from kwai.api.schemas.resources import ( 

6 TeamResourceIdentifier, 

7 TrainingDefinitionResourceIdentifier, 

8) 

9from kwai.api.v1.trainings.schemas.team import TeamAttributes, TeamResource 

10from kwai.core.json_api import Document, Relationship, ResourceData, ResourceMeta 

11from kwai.modules.training.trainings.training_definition import TrainingDefinitionEntity 

12 

13 

14class TrainingDefinitionAttributes(BaseModel): 

15 """Attributes for a training definition JSON:API resource.""" 

16 

17 name: str 

18 description: str 

19 weekday: int 

20 start_time: str 

21 end_time: str 

22 timezone: str 

23 active: bool 

24 location: str 

25 remark: str | None = None 

26 

27 

28class TrainingDefinitionRelationships(BaseModel): 

29 """Relationships for a training definition JSON:API resource.""" 

30 

31 team: Relationship[TeamResourceIdentifier] | None = None 

32 

33 

34class TrainingDefinitionResource( 

35 TrainingDefinitionResourceIdentifier, 

36 ResourceData[ 

37 TrainingDefinitionAttributes, 

38 TrainingDefinitionRelationships, 

39 ], 

40): 

41 """A JSON:API resource for a training definition.""" 

42 

43 

44class TrainingDefinitionDocument(Document[TrainingDefinitionResource, TeamResource]): 

45 """A JSON:API document for one or more training definitions.""" 

46 

47 @classmethod 

48 def create( 

49 cls, training_definition: TrainingDefinitionEntity 

50 ) -> "TrainingDefinitionDocument": 

51 """Create a document from a training definition entity.""" 

52 data = TrainingDefinitionResource( 

53 id=str(training_definition.id), 

54 meta=ResourceMeta( 

55 created_at=str(training_definition.traceable_time.created_at), 

56 updated_at=str(training_definition.traceable_time.updated_at), 

57 ), 

58 attributes=TrainingDefinitionAttributes( 

59 name=training_definition.name, 

60 description=training_definition.description, 

61 weekday=training_definition.weekday.value, 

62 start_time=str(training_definition.period.start), 

63 end_time=str(training_definition.period.end), 

64 timezone=training_definition.period.timezone, 

65 active=training_definition.active, 

66 location=training_definition.location, 

67 remark=training_definition.remark or "", 

68 ), 

69 ) 

70 

71 included: set[TeamResource] = set() 

72 if training_definition.team is not None: 

73 data.relationships = TrainingDefinitionRelationships( 

74 team=Relationship[TeamResourceIdentifier]( 

75 data=TeamResourceIdentifier(id=str(training_definition.team.id)) 

76 ) 

77 ) 

78 included.add( 

79 TeamResource( 

80 id=str(training_definition.team.id), 

81 attributes=TeamAttributes(name=str(training_definition.team.name)), 

82 ) 

83 ) 

84 

85 return TrainingDefinitionDocument(data=data, included=included)