Coverage for src/tests/fixtures/training/training_schemas.py: 100%

1"""Module for defining fixture factories for training schedules.""" 

2 

3import pytest 

4 

5from kwai_bc_training.coaches.coach import CoachEntity 

6from kwai_bc_training.trainings.training_schedule import TrainingScheduleEntity 

7from kwai_bc_training.trainings.training_schedule_db_repository import ( 

8 TrainingScheduleDbRepository, 

9) 

10from kwai_core.db.database import Database 

11from kwai_core.db.uow import UnitOfWork 

12from kwai_core.domain.value_objects.owner import Owner 

13from kwai_core.domain.value_objects.time_period import TimePeriod 

14from kwai_core.domain.value_objects.weekday import Weekday 

15 

16 

17@pytest.fixture 

18def make_training_schedule(owner: Owner, make_coach): 

19 """A factory fixture for creating a training schedule.""" 

20 

21 def _make_training_schedule( 

22 name: str | None = None, 

23 description: str | None = None, 

24 weekday: Weekday | None = None, 

25 period: TimePeriod | None = None, 

26 coach: CoachEntity | None = None, 

27 ) -> TrainingScheduleEntity: 

28 coach = coach or make_coach() 

29 return TrainingScheduleEntity( 

30 name=name or "A training", 

31 description=description or "A training schedule", 

32 weekday=weekday or Weekday.MONDAY, 

33 period=period 

34 or TimePeriod.create_from_string("18:00", "19:00", "Europe/Brussels"), 

35 owner=owner, 

36 coaches=frozenset({coach}), 

37 ) 

38 

39 return _make_training_schedule 

40 

41 

42@pytest.fixture 

43async def make_training_schedule_in_db( 

44 database: Database, make_training_schedule, make_coach_in_db 

45): 

46 """A fixture factory for a training schedule in the database.""" 

47 created_entities = [] 

48 

49 async def _make_training_schedule_in_db( 

50 training_schedule: TrainingScheduleEntity | None = None, 

51 coach: CoachEntity | None = None, 

52 ) -> TrainingScheduleEntity: 

53 coach = coach or await make_coach_in_db() 

54 training_schedule = training_schedule or make_training_schedule(coach=coach) 

55 created_entities.append(training_schedule) 

56 repo = TrainingScheduleDbRepository(database) 

57 async with UnitOfWork(database): 

58 training_schedule = await repo.create(training_schedule) 

59 created_entities.append(training_schedule) 

60 

61 return training_schedule 

62 

63 yield _make_training_schedule_in_db 

64 

65 repo = TrainingScheduleDbRepository(database) 

66 for entity in created_entities: 

67 async with UnitOfWork(database): 

68 await repo.delete(entity)