Coverage for sm/utils_starterpack.py: 89%

1from pathlib import Path 

2import yaml 

3from django.conf import settings 

4from django.contrib.auth.models import Group 

5from django.db import transaction 

6from vendor.models import Model as Vendor 

7from operatingsystem.models import Model as OS 

8from typing import Dict, List, Any 

9 

10 

11def _load_fixture(app_label: str) -> List[Any]: 

12 """Loads a YAML fixture from an app's fixtures dir; [] if absent.""" 

13 path = Path(settings.BASE_DIR) / app_label / "fixtures" / "01_initial.yaml" 

14 if not path.is_file(): 

15 return [] 

16 with open(path) as f: 

17 return yaml.safe_load(f) or [] 

18 

19 

20def import_starter_pack(group: Group) -> Dict[str, int]: 

21 """ 

22 Imports vendors and operating systems from fixtures into the specified 

23 group. Runs atomically so a partial import can never leave the group in a 

24 half-populated state. 

25 """ 

26 # Vendor name (as referenced from the OS fixture) -> group-scoped Vendor 

27 vendor_map = {} 

28 vendor_count = 0 

29 os_count = 0 

30 

31 with transaction.atomic(): 

32 vendor_data = _load_fixture("vendor") 

33 for item in vendor_data: 

34 if not isinstance(item, dict) or item.get("model") != "vendor.model": 

35 continue 

36 fields = item["fields"] 

37 vendor, created = Vendor.objects.get_or_create( 

38 name=fields["name"], 

39 group=group, 

40 defaults={ 

41 "is_hardware": fields.get("is_hardware", True), 

42 "is_software": fields.get("is_software", True), 

43 }, 

44 ) 

45 vendor_map[fields["name"]] = vendor 

46 if created: 

47 vendor_count += 1 

48 

49 os_data = _load_fixture("operatingsystem") 

50 for item in os_data: 

51 if ( 

52 not isinstance(item, dict) 

53 or item.get("model") != "operatingsystem.model" 

54 ): 

55 continue 

56 fields = item["fields"] 

57 vendor_names = fields.get("vendor") 

58 if isinstance(vendor_names, list): 

59 vendor_name = vendor_names[0] if vendor_names else None 

60 else: 

61 vendor_name = vendor_names 

62 

63 vendor = vendor_map.get(vendor_name) if vendor_name else None 

64 if vendor is None: 

65 continue 

66 _, created = OS.objects.get_or_create( 

67 version=fields["version"], vendor=vendor, group=group 

68 ) 

69 if created: 

70 os_count += 1 

71 

72 return {"vendors": vendor_count, "os": os_count}