Coverage for testing/factories.py: 73%

1""" 

2crate_anon/testing/factories.py 

3 

4=============================================================================== 

5 

6 Copyright (C) 2015, University of Cambridge, Department of Psychiatry. 

7 Created by Rudolf Cardinal (rnc1001@cam.ac.uk). 

8 

9 This file is part of CRATE. 

10 

11 CRATE is free software: you can redistribute it and/or modify 

12 it under the terms of the GNU General Public License as published by 

13 the Free Software Foundation, either version 3 of the License, or 

14 (at your option) any later version. 

15 

16 CRATE is distributed in the hope that it will be useful, 

17 but WITHOUT ANY WARRANTY; without even the implied warranty of 

18 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

19 GNU General Public License for more details. 

20 

21 You should have received a copy of the GNU General Public License 

22 along with CRATE. If not, see <https://www.gnu.org/licenses/>. 

23 

24=============================================================================== 

25 

26**Factory Boy SQL Alchemy test factories.** 

27 

28""" 

29 

30import random 

31from typing import TYPE_CHECKING 

32 

33from cardinal_pythonlib.classes import all_subclasses 

34import factory 

35import factory.random 

36from faker import Faker 

37 

38from crate_anon.testing.models import EnumColours, FilenameDoc, Note, Patient 

39from crate_anon.testing.providers import register_all_providers 

40 

41if TYPE_CHECKING: 

42 from factory.builder import Resolver 

43 from sqlalchemy.orm.session import Session 

44 

45 

46# When running with pytest sqlalchemy_session gets poked in by 

47# DatabaseTestCase.setUp(). Otherwise call 

48# set_sqlalchemy_session_on_all_factories() 

49class AnonTestBaseFactory(factory.alchemy.SQLAlchemyModelFactory): 

50 pass 

51 

52 

53class SecretBaseFactory(factory.alchemy.SQLAlchemyModelFactory): 

54 pass 

55 

56 

57class SourceTestBaseFactory(factory.alchemy.SQLAlchemyModelFactory): 

58 pass 

59 

60 

61def set_sqlalchemy_session_on_all_factories( 

62 factory_base_class: factory.alchemy.SQLAlchemyModelFactory, 

63 dbsession: "Session", 

64) -> None: 

65 for factory_class in all_subclasses(factory_base_class): 

66 factory_class._meta.sqlalchemy_session = dbsession 

67 

68 

69# ============================================================================= 

70# Randomness 

71# ============================================================================= 

72 

73 

74def coin(p: float = 0.5) -> bool: 

75 """ 

76 Biased coin toss. Returns ``True`` with probability ``p``. 

77 """ 

78 return random.random() < p 

79 

80 

81class Fake: 

82 # MB 2024-02-19 

83 # Factory Boy has its own interface to Faker (factory.Faker()). This 

84 # takes a function to be called at object generation time and as far as I 

85 # can tell this doesn't support being able to create fake data based on 

86 # other fake attributes such as notes for a patient. You can work 

87 # around this by adding a lot of logic to the factories. To me it makes 

88 # sense to keep the factories simple and do as much as possible of the 

89 # content generation in the providers. So we call Faker directly instead. 

90 en_gb = Faker("en_GB") # For UK postcodes, phone numbers etc 

91 en_us = Faker("en_US") # en_GB gives Lorem ipsum for pad words. 

92 

93 

94register_all_providers(Fake.en_gb) 

95 

96 

97class DemoFactory(SourceTestBaseFactory): 

98 class Meta: 

99 abstract = True 

100 

101 

102class DemoPatientFactory(DemoFactory): 

103 class Meta: 

104 model = Patient 

105 

106 patient_id = factory.Sequence(lambda n: n + 1) 

107 

108 sex = factory.LazyFunction(Fake.en_gb.sex) 

109 

110 @factory.lazy_attribute 

111 def forename(obj: "Resolver") -> str: 

112 return Fake.en_gb.forename(obj.sex) 

113 

114 surname = factory.LazyFunction(Fake.en_gb.last_name) 

115 dob = factory.LazyFunction(Fake.en_gb.consistent_date_of_birth) 

116 nhsnum = factory.LazyFunction(Fake.en_gb.nhs_number) 

117 

118 phone = factory.LazyFunction(Fake.en_gb.phone_number) 

119 postcode = factory.LazyFunction(Fake.en_gb.postcode) 

120 

121 @factory.lazy_attribute 

122 def related_patient(obj: "Resolver") -> int: 

123 if obj.patient_id == 1: 

124 return None 

125 

126 related_patient_id = obj.patient_id - 1 

127 session = DemoPatientFactory._meta.sqlalchemy_session 

128 related_patient = ( 

129 session.query(Patient) 

130 .filter(Patient.patient_id == related_patient_id) 

131 .first() 

132 ) 

133 

134 return related_patient 

135 

136 related_patient_relationship = factory.LazyFunction( 

137 Fake.en_gb.relationship 

138 ) 

139 

140 @factory.lazy_attribute 

141 def colour(obj: "Resolver") -> EnumColours: 

142 return EnumColours.blue if coin() else None 

143 

144 @factory.post_generation 

145 def notes(obj: "Resolver", create: bool, extracted: int, **kwargs) -> None: 

146 if not create: 

147 return 

148 

149 if extracted: 

150 DemoNoteFactory.create_batch(size=extracted, patient=obj, **kwargs) 

151 

152 

153class DemoNoteFactory(DemoFactory): 

154 class Meta: 

155 model = Note 

156 

157 class Params: 

158 words_per_note = 100 

159 

160 note_datetime = factory.LazyFunction(Fake.en_gb.incrementing_date) 

161 

162 @factory.lazy_attribute 

163 def note(obj: "Resolver") -> str: 

164 # Use en_US because you get Lorem ipsum with en_GB. 

165 pad_paragraph = Fake.en_us.paragraph( 

166 nb_sentences=obj.words_per_note / 2, # way more than we need 

167 ) 

168 

169 return Fake.en_gb.patient_note( 

170 forename=obj.patient.forename, 

171 surname=obj.patient.surname, 

172 sex=obj.patient.sex, 

173 dob=obj.patient.dob, 

174 nhs_number=obj.patient.nhsnum, 

175 patient_id=obj.patient.patient_id, 

176 note_datetime=obj.note_datetime, 

177 relation_name=obj.patient.related_patient_name, 

178 relation_relationship=obj.patient.related_patient_relationship, 

179 words_per_note=obj.words_per_note, 

180 pad_paragraph=pad_paragraph, 

181 ) 

182 

183 

184class DemoFilenameDocFactory(DemoFactory): 

185 class Meta: 

186 model = FilenameDoc 

187 

188 file_datetime = factory.LazyFunction(Fake.en_gb.incrementing_date) 

189 

190 @factory.lazy_attribute 

191 def filename(obj: "Resolver") -> str: 

192 # Use en_US because you get Lorem ipsum with en_GB. 

193 pad_paragraph = Fake.en_us.paragraph(nb_sentences=50) 

194 

195 return Fake.en_gb.patient_filename( 

196 forename=obj.patient.forename, 

197 surname=obj.patient.surname, 

198 sex=obj.patient.sex, 

199 dob=obj.patient.dob, 

200 nhs_number=obj.patient.nhsnum, 

201 patient_id=obj.patient.patient_id, 

202 pad_paragraph=pad_paragraph, 

203 )