Set level in new fake incidents

src/argus/dev/management/commands/create_fake_incident.py

@@ -1,21 +1,50 @@
+import argparse
+from random import randint
+
 from django.core.management.base import BaseCommand
 
+from argus.incident.constants import MIN_INCIDENT_LEVEL, MAX_INCIDENT_LEVEL
 from argus.incident.models import create_fake_incident
 
 
+class Range(argparse.Action):
+    def __init__(self, minimum=None, maximum=None, *args, **kwargs):
+        self.min = minimum
+        self.max = maximum
+        kwargs["metavar"] = "[%d-%d]" % (self.min, self.max)
+        super(Range, self).__init__(*args, **kwargs)
+
+    def __call__(self, parser, namespace, value, option_string=None):
+        if not (self.min <= value <= self.max):
+            msg = "invalid choice: %r (choose from [%d-%d])" % (value, self.min, self.max)
+            raise argparse.ArgumentError(self, msg)
+        setattr(namespace, self.dest, value)
+
+
 class Command(BaseCommand):
     help = "Create fake Incident"
 
     def add_arguments(self, parser):
-        parser.add_argument("-t", "--tags", nargs="+", type=str, help="Add the listed tags to the incident")
         parser.add_argument("-b", "--batch-size", type=int, help="Create <batch size> incidents in one go")
         parser.add_argument("-d", "--description", type=str, help="Use this description for the incident")
+        parser.add_argument(
+            "-l",
+            "--level",
+            type=int,
+            action=Range,
+            minimum=MIN_INCIDENT_LEVEL,
+            maximum=MAX_INCIDENT_LEVEL,
+            default=0,
+            help="Set level to <level>, otherwise a random number within the correct range is used",
+        )
+        parser.add_argument("-t", "--tags", nargs="+", type=str, help="Add the listed tags to the incident")
         parser.add_argument("--stateless", action="store_true", help="Create a stateless incident (end_time = None)")
 
     def handle(self, *args, **options):
         tags = options.get("tags") or []
         description = options.get("description") or None
         batch_size = options.get("batch_size") or 1
+        level = options.get("level") or randint(MIN_INCIDENT_LEVEL, MAX_INCIDENT_LEVEL)
         stateful = False if options.get("stateless") else True
         for i in range(batch_size):
             create_fake_incident(tags=tags, description=description, stateful=stateful)

src/argus/incident/constants.py

@@ -1,5 +1,7 @@
 # Prevent import loops
 # DO NOT import anything here, ever
 
-INCIDENT_LEVELS = (1, 2, 3, 4, 5)
+MIN_INCIDENT_LEVEL = 1  # Do not override
+MAX_INCIDENT_LEVEL = 5
+INCIDENT_LEVELS = tuple(range(MIN_INCIDENT_LEVEL, MAX_INCIDENT_LEVEL + 1))
 INCIDENT_LEVEL_CHOICES = zip(INCIDENT_LEVELS, map(str, INCIDENT_LEVELS))

src/argus/incident/models.py

@@ -13,7 +13,7 @@
 
 from argus.auth.models import User
 from argus.util.datetime_utils import INFINITY_REPR, get_infinity_repr
-from .constants import INCIDENT_LEVELS, INCIDENT_LEVEL_CHOICES
+from .constants import INCIDENT_LEVELS, INCIDENT_LEVEL_CHOICES, MIN_INCIDENT_LEVEL, MAX_INCIDENT_LEVEL
 from .fields import DateTimeInfinityField
 from .validators import validate_lowercase, validate_key
 
@@ -28,7 +28,7 @@ def get_or_create_default_instances():
     return (argus_user, sst, ss)
 
 
-def create_fake_incident(tags=None, description=None, stateful=True):
+def create_fake_incident(tags=None, description=None, stateful=True, level=None):
     argus_user, _, source_system = get_or_create_default_instances()
     end_time = INFINITY_REPR if stateful else None
 
@@ -47,6 +47,7 @@ def create_fake_incident(tags=None, description=None, stateful=True):
         source_incident_id=source_incident_id,
         source=source_system,
         description=description,
+        level=level or randint(MIN_INCIDENT_LEVEL, MAX_INCIDENT_LEVEL),
     )
 
     taglist = [("location", "argus"), ("object", f"{incident.id}"), ("problem_type", "test")]