Coverage for sm / utils_starterpack.py: 22%

1import yaml 

2import os 

3from django.conf import settings 

4from django.contrib.auth.models import Group 

5from vendor.models import Model as Vendor 

6from operatingsystem.models import Model as OS 

7from typing import Dict, List, Any, Optional, Union 

8 

9 

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

11 """ 

12 Imports vendors and operating systems from fixtures into the specified group. 

13 """ 

14 vendor_count = 0 

15 os_count = 0 

16 

17 # 1. Import Vendors 

18 vendor_fixture_path = os.path.join( 

19 settings.BASE_DIR, "vendor", "fixtures", "01_initial.yaml" 

20 ) 

21 with open(vendor_fixture_path) as f: 

22 vendor_data = yaml.safe_load(f) 

23 

24 vendor_map = {} # To keep track of created vendors for OS lookup 

25 

26 for item in vendor_data: 

27 if item["model"] == "vendor.model": 

28 fields = item["fields"] 

29 vendor, created = Vendor.objects.get_or_create( 

30 name=fields["name"], 

31 group=group, 

32 defaults={ 

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

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

35 }, 

36 ) 

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

38 if created: 

39 vendor_count += 1 

40 

41 # 2. Import Operating Systems 

42 os_fixture_path = os.path.join( 

43 settings.BASE_DIR, "operatingsystem", "fixtures", "01_initial.yaml" 

44 ) 

45 with open(os_fixture_path) as f: 

46 os_data = yaml.safe_load(f) 

47 

48 for item in os_data: 

49 if item["model"] == "operatingsystem.model": 

50 fields = item["fields"] 

51 vendor_names = fields["vendor"] 

52 if isinstance(vendor_names, list): 

53 vendor_name = vendor_names[0] 

54 else: 

55 vendor_name = vendor_names 

56 

57 vendor = vendor_map.get(vendor_name) 

58 if vendor: 

59 os_obj, created = OS.objects.get_or_create( 

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

61 ) 

62 if created: 

63 os_count += 1 

64 

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