Coverage for utilities/testing/utils.py: 0%

1import json 

2import logging 

3import random 

4import re 

5import string 

6from contextlib import contextmanager 

7 

8from django.contrib.auth.models import Permission 

9from django.utils.text import slugify 

10 

11from core.models import ObjectType 

12from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Site 

13from extras.choices import CustomFieldTypeChoices 

14from extras.models import CustomField, Tag 

15from ipam.models import IPAddress 

16from users.models import User 

17from virtualization.models import Cluster, ClusterType, VirtualMachine 

18 

19 

20def post_data(data): 

21 """ 

22 Take a dictionary of test data (suitable for comparison to an instance) and return a dict suitable for POSTing. 

23 """ 

24 ret = {} 

25 

26 for key, value in data.items(): 

27 if value is None: 

28 ret[key] = '' 

29 elif type(value) in (list, tuple): 

30 if value and hasattr(value[0], 'pk'): 

31 # Value is a list of instances 

32 ret[key] = [v.pk for v in value] 

33 else: 

34 ret[key] = value 

35 elif hasattr(value, 'pk'): 

36 # Value is an instance 

37 ret[key] = value.pk 

38 else: 

39 ret[key] = str(value) 

40 

41 return ret 

42 

43 

44def create_test_device(name, site=None, **attrs): 

45 """ 

46 Convenience method for creating a Device (e.g. for component testing). 

47 """ 

48 if site is None: 

49 site, _ = Site.objects.get_or_create(name='Site 1', slug='site-1') 

50 manufacturer, _ = Manufacturer.objects.get_or_create(name='Manufacturer 1', slug='manufacturer-1') 

51 devicetype, _ = DeviceType.objects.get_or_create(model='Device Type 1', manufacturer=manufacturer) 

52 devicerole, _ = DeviceRole.objects.get_or_create(name='Device Role 1', slug='device-role-1') 

53 device = Device.objects.create(name=name, site=site, device_type=devicetype, role=devicerole, **attrs) 

54 

55 return device 

56 

57 

58def create_test_virtualmachine(name): 

59 """ 

60 Convenience method for creating a VirtualMachine. 

61 """ 

62 cluster_type, _ = ClusterType.objects.get_or_create(name='Cluster Type 1', slug='cluster-type-1') 

63 cluster, _ = Cluster.objects.get_or_create(name='Cluster 1', type=cluster_type) 

64 virtual_machine = VirtualMachine.objects.create(name=name, cluster=cluster) 

65 

66 return virtual_machine 

67 

68 

69def create_test_nat_ip_pair( 

70 real_address='10.0.0.10/32', nat_address='198.51.100.10/32', inside_interface=None, outside_interface=None 

71): 

72 """ 

73 Convenience method for creating an inside IP and its NAT outside IP. 

74 

75 Optionally, assign either address to an Interface or VMInterface. 

76 Returns (real_ip, nat_ip). 

77 """ 

78 real_ip = IPAddress(address=real_address) 

79 if inside_interface is not None: 

80 real_ip.assigned_object = inside_interface 

81 real_ip.save() 

82 

83 nat_ip = IPAddress(address=nat_address, nat_inside=real_ip) 

84 if outside_interface is not None: 

85 nat_ip.assigned_object = outside_interface 

86 nat_ip.save() 

87 

88 return real_ip, nat_ip 

89 

90 

91def create_test_user(username='testuser', permissions=None): 

92 """ 

93 Create a User with the given permissions. 

94 """ 

95 user = User.objects.create_user(username=username) 

96 if permissions is None: 

97 permissions = () 

98 for perm_name in permissions: 

99 app, codename = perm_name.split('.') 

100 perm = Permission.objects.get(content_type__app_label=app, codename=codename) 

101 user.user_permissions.add(perm) 

102 

103 return user 

104 

105 

106def create_tags(*names): 

107 """ 

108 Create and return a Tag instance for each name given. 

109 """ 

110 tags = [Tag(name=name, slug=slugify(name)) for name in names] 

111 Tag.objects.bulk_create(tags) 

112 return tags 

113 

114 

115def extract_form_failures(content): 

116 """ 

117 Given raw HTML content from an HTTP response, return a list of form errors. 

118 """ 

119 FORM_ERROR_REGEX = r'<!-- FORM-ERROR (.*) -->' 

120 return re.findall(FORM_ERROR_REGEX, str(content)) 

121 

122 

123@contextmanager 

124def disable_warnings(logger_name): 

125 """ 

126 Temporarily suppress expected warning messages to keep the test output clean. 

127 """ 

128 logger = logging.getLogger(logger_name) 

129 current_level = logger.level 

130 logger.setLevel(logging.ERROR) 

131 yield 

132 logger.setLevel(current_level) 

133 

134 

135@contextmanager 

136def disable_logging(level=logging.CRITICAL): 

137 """ 

138 Temporarily suppress log messages at or below the specified level (default: critical). 

139 """ 

140 # Capture the current disable level so it can be restored on exit (rather than assuming 

141 # NOTSET), which keeps nested calls well-behaved. The teardown runs inside a finally block so 

142 # logging is always restored even if the wrapped block raises. 

143 previous_level = logging.root.manager.disable 

144 logging.disable(level) 

145 try: 

146 yield 

147 finally: 

148 logging.disable(previous_level) 

149 

150 

151# 

152# Custom field testing 

153# 

154 

155DUMMY_CF_DATA = { 

156 'text_field': 'foo123', 

157 'integer_field': 456, 

158 'decimal_field': 456.12, 

159 'boolean_field': True, 

160 'json_field': {'abc': 123}, 

161} 

162 

163 

164def add_custom_field_data(form_data, model): 

165 """ 

166 Create some custom fields for the model and add a value for each to the form data. 

167 

168 Args: 

169 form_data: The dictionary of form data to be updated 

170 model: The model of the object the form seeks to create or modify 

171 """ 

172 object_type = ObjectType.objects.get_for_model(model) 

173 custom_fields = ( 

174 CustomField(type=CustomFieldTypeChoices.TYPE_TEXT, name='text_field', default='foo'), 

175 CustomField(type=CustomFieldTypeChoices.TYPE_INTEGER, name='integer_field', default=123), 

176 CustomField(type=CustomFieldTypeChoices.TYPE_DECIMAL, name='decimal_field', default=123.45), 

177 CustomField(type=CustomFieldTypeChoices.TYPE_BOOLEAN, name='boolean_field', default=False), 

178 CustomField(type=CustomFieldTypeChoices.TYPE_JSON, name='json_field', default='{"x": "y"}'), 

179 ) 

180 CustomField.objects.bulk_create(custom_fields) 

181 for cf in custom_fields: 

182 cf.object_types.set([object_type]) 

183 

184 form_data.update({ 

185 f'cf_{k}': v if type(v) is str else json.dumps(v) 

186 for k, v in DUMMY_CF_DATA.items() 

187 }) 

188 

189 

190# 

191# Misc utilities 

192# 

193 

194def get_random_string(length, charset=None): 

195 """ 

196 Return a random string of the given length. 

197 """ 

198 characters = string.ascii_letters + string.digits # a-z, A-Z, 0-9 

199 return ''.join(random.choice(characters) for __ in range(length)) 

200 

201 

202# 

203# Database routing 

204# 

205 

206class UnpinnedQuery(Exception): 

207 """Raised when a query which should have been pinned to a connection is routed instead.""" 

208 

209 

210class PinnedConnectionRouter: 

211 """ 

212 Fails any read or write of the given models which is not pinned to an explicit database 

213 alias. Django consults DATABASE_ROUTERS only for queries which name no connection, so a 

214 signal handler which threads through the alias supplied by the signal never reaches 

215 this router. Each test leaves out the model being saved, as Django routes that save 

216 itself. 

217 """ 

218 def __init__(self, *models): 

219 self.models = models 

220 

221 def _check(self, model, **hints): 

222 if model in self.models: 

223 raise UnpinnedQuery(f"{model.__name__} query was routed rather than pinned to a connection") 

224 

225 db_for_read = _check 

226 db_for_write = _check