Coverage for dcim/management/commands/trace_paths.py: 46%

1from django.core.management.base import BaseCommand 

2from django.core.management.color import no_style 

3from django.db import connection 

4from django.db.models import Q 

5 

6from dcim.models import CablePath, ConsolePort, ConsoleServerPort, Interface, PowerFeed, PowerOutlet, PowerPort 

7from dcim.signals import create_cablepaths 

8 

9ENDPOINT_MODELS = ( 

10 ConsolePort, 

11 ConsoleServerPort, 

12 Interface, 

13 PowerFeed, 

14 PowerOutlet, 

15 PowerPort 

16) 

17 

18 

19class Command(BaseCommand): 

20 help = "Generate any missing cable paths among all cable termination objects in NetBox" 

21 

22 def add_arguments(self, parser): 

23 parser.add_argument( 

24 "--force", action='store_true', dest='force', 

25 help="Force recalculation of all existing cable paths" 

26 ) 

27 parser.add_argument( 

28 "--no-input", action='store_true', dest='no_input', 

29 help="Do not prompt user for any input/confirmation" 

30 ) 

31 

32 def draw_progress_bar(self, percentage): 

33 """ 

34 Draw a simple progress bar 20 increments wide illustrating the specified percentage. 

35 """ 

36 bar_size = int(percentage / 5) 

37 self.stdout.write(f"\r [{'#' * bar_size}{' ' * (20 - bar_size)}] {int(percentage)}%", ending='') 

38 

39 def handle(self, *model_names, **options): 

40 

41 # If --force was passed, first delete all existing CablePaths 

42 if options['force']: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true

43 cable_paths = CablePath.objects.all() 

44 paths_count = cable_paths.count() 

45 

46 # Prompt the user to confirm recalculation of all paths 

47 if paths_count and not options['no_input']: 

48 self.stdout.write(self.style.ERROR("WARNING: Forcing recalculation of all cable paths.")) 

49 self.stdout.write( 

50 f"This will delete and recalculate all {paths_count} existing cable paths. Are you sure?" 

51 ) 

52 confirmation = input("Type yes to confirm: ") 

53 if confirmation != 'yes': 

54 self.stdout.write(self.style.SUCCESS("Aborting")) 

55 return 

56 

57 # Delete all existing CablePath instances 

58 self.stdout.write(f"Deleting {paths_count} existing cable paths...") 

59 deleted_count, _ = CablePath.objects.all().delete() 

60 self.stdout.write(self.style.SUCCESS(f' Deleted {deleted_count} paths')) 

61 

62 # Reinitialize the model's PK sequence 

63 self.stdout.write('Resetting database sequence for CablePath model') 

64 sequence_sql = connection.ops.sequence_reset_sql(no_style(), [CablePath]) 

65 with connection.cursor() as cursor: 

66 for sql in sequence_sql: 

67 cursor.execute(sql) 

68 

69 # Retrace paths 

70 for model in ENDPOINT_MODELS: 

71 params = Q(cable__isnull=False) 

72 if hasattr(model, 'wireless_link'): 

73 params |= Q(wireless_link__isnull=False) 

74 origins = model.objects.filter(params) 

75 if not options['force']: 75 ↛ 77line 75 didn't jump to line 77 because the condition on line 75 was always true

76 origins = origins.filter(_path__isnull=True) 

77 origins_count = origins.count() 

78 if not origins_count: 78 ↛ 81line 78 didn't jump to line 81 because the condition on line 78 was always true

79 self.stdout.write(f'Found no missing {model._meta.verbose_name} paths; skipping') 

80 continue 

81 self.stdout.write(f'Retracing {origins_count} cabled {model._meta.verbose_name_plural}...') 

82 i = 0 

83 for i, obj in enumerate(origins, start=1): 

84 create_cablepaths([obj]) 

85 if not i % 100: 

86 self.draw_progress_bar(i * 100 / origins_count) 

87 self.draw_progress_bar(100) 

88 self.stdout.write(self.style.SUCCESS(f'\n Retraced {i} {model._meta.verbose_name_plural}')) 

89 

90 self.stdout.write(self.style.SUCCESS('Finished.'))