Coverage for dcim/management/commands/buildschema.py: 0%

1import json 

2import os 

3 

4from django.conf import settings 

5from django.core.management.base import BaseCommand 

6from jinja2 import Environment, FileSystemLoader 

7 

8from dcim.choices import * 

9from netbox.choices import WeightUnitChoices 

10from wireless.choices import WirelessRoleChoices 

11 

12TEMPLATE_FILENAME = 'devicetype_schema.jinja2' 

13OUTPUT_FILENAME = 'contrib/generated_schema.json' 

14 

15CHOICES_MAP = { 

16 'airflow_choices': DeviceAirflowChoices, 

17 'weight_unit_choices': WeightUnitChoices, 

18 'subdevice_role_choices': SubdeviceRoleChoices, 

19 'console_port_type_choices': ConsolePortTypeChoices, 

20 'console_server_port_type_choices': ConsolePortTypeChoices, 

21 'power_port_type_choices': PowerPortTypeChoices, 

22 'power_outlet_type_choices': PowerOutletTypeChoices, 

23 'power_outlet_feedleg_choices': PowerOutletFeedLegChoices, 

24 'interface_type_choices': InterfaceTypeChoices, 

25 'interface_poe_mode_choices': InterfacePoEModeChoices, 

26 'interface_poe_type_choices': InterfacePoETypeChoices, 

27 'interface_rf_role_choices': WirelessRoleChoices, 

28 'front_port_type_choices': PortTypeChoices, 

29 'rear_port_type_choices': PortTypeChoices, 

30} 

31 

32 

33class Command(BaseCommand): 

34 help = "Generate JSON schema for validating NetBox device type definitions" 

35 

36 def add_arguments(self, parser): 

37 parser.add_argument( 

38 '--write', 

39 action='store_true', 

40 help="Write the generated schema to file" 

41 ) 

42 

43 def handle(self, *args, **kwargs): 

44 # Initialize template 

45 template_loader = FileSystemLoader(searchpath=f'{settings.TEMPLATES_DIR}/extras/schema/') 

46 template_env = Environment(loader=template_loader) 

47 template = template_env.get_template(TEMPLATE_FILENAME) 

48 

49 # Render template 

50 context = { 

51 key: json.dumps(choices.values()) 

52 for key, choices in CHOICES_MAP.items() 

53 } 

54 rendered = template.render(**context) 

55 

56 if kwargs['write']: 

57 # $root/contrib/generated_schema.json 

58 filename = os.path.join(os.path.split(settings.BASE_DIR)[0], OUTPUT_FILENAME) 

59 with open(filename, mode='w', encoding='UTF-8') as f: 

60 f.write(json.dumps(json.loads(rendered), indent=4)) 

61 f.write('\n') 

62 f.close() 

63 self.stdout.write(self.style.SUCCESS(f"Schema written to {filename}.")) 

64 else: 

65 self.stdout.write(rendered)