Add source argument to create fake incident

Author: johannaengland

changelog.d/+source_fake_incident.added.md

@@ -0,0 +1 @@
+Add source argument to create_fake_incident CLI command

docs/development/management-commands.rst

@@ -94,6 +94,13 @@ To add a custom description add the `-d` flag to the command as such:
 
         $ python manage.py create_fake_incident -d "Custom description"
 
+To use a different source than 'argus' add the `-s` flag to the command as
+such:
+
+    .. code:: console
+
+        $ python manage.py create_fake_incident -s "Notargus"
+
 To set the level of the incident add the `-l` flag to the command as such
 and choose a level between 1 and 5 (1 being the highest severity, 5 the
 lowest):

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

@@ -2,7 +2,9 @@
 import json
 from pathlib import Path
 
+from django.core.management import CommandError
 from django.core.management.base import BaseCommand
+from django.core.exceptions import ValidationError
 
 from argus.incident.constants import Level
 from argus.incident.models import create_fake_incident
@@ -39,6 +41,12 @@ def add_arguments(self, parser):
             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(
+            "-s",
+            "--source",
+            type=str,
+            help="Use this source for the incident (the source needs to exist, see 'create_source' for creating one)",
+        )
         parser.add_argument("--stateless", action="store_true", help="Create a stateless incident (end_time = None)")
         metadata_parser = parser.add_mutually_exclusive_group()
         metadata_parser.add_argument(
@@ -55,6 +63,7 @@ def add_arguments(self, parser):
     def handle(self, *args, **options):
         tags = options.get("tags") or []
         description = options.get("description") or None
+        source = options.get("source") or None
         batch_size = options.get("batch_size") or 1
         level = options.get("level") or None
         stateful = False if options.get("stateless") else True
@@ -65,10 +74,14 @@ def handle(self, *args, **options):
                 with metadata_path.open() as jsonfile:
                     metadata = json.load(jsonfile)
         for i in range(batch_size):
-            create_fake_incident(
-                tags=tags,
-                description=description,
-                stateful=stateful,
-                level=level,
-                metadata=metadata,
-            )
+            try:
+                create_fake_incident(
+                    tags=tags,
+                    description=description,
+                    source=source,
+                    stateful=stateful,
+                    level=level,
+                    metadata=metadata,
+                )
+            except (ValueError, ValidationError) as e:
+                raise CommandError(str(e))

src/argus/incident/models.py

@@ -30,8 +30,13 @@ def get_or_create_default_instances():
     return (argus_user, sst, ss)
 
 
-def create_fake_incident(tags=None, description=None, stateful=True, level=None, metadata=None):
+def create_fake_incident(tags=None, description=None, source=None, stateful=True, level=None, metadata=None):
     argus_user, _, source_system = get_or_create_default_instances()
+    if source:
+        try:
+            source_system = SourceSystem.objects.get(name=source)
+        except SourceSystem.DoesNotExist:
+            raise ValueError(f"No source with the name '{source}' exists.")
     end_time = INFINITY_REPR if stateful else None
 
     MAX_ID = 2**32 - 1

tests/dev/test_create_fake_incident.py

@@ -3,6 +3,7 @@
 from django.core.management import CommandError, call_command
 from django.test import TestCase
 
+from argus.incident.factories import SourceSystemFactory, SourceSystemTypeFactory, SourceUserFactory
 from argus.incident.models import Incident, Tag
 from argus.util.testing import connect_signals, disconnect_signals
 
@@ -67,6 +68,26 @@ def test_create_fake_incident_will_create_single_fake_incident_with_set_descript
             .exists()
         )
 
+    def test_create_fake_incident_will_create_single_fake_incident_with_set_source(self):
+        previous_incidents_pks = [incident.id for incident in Incident.objects.all()]
+
+        source_name = "notargus"
+        sst = SourceSystemTypeFactory(name="source_type")
+        user = SourceUserFactory(username=source_name)
+        SourceSystemFactory(name=source_name, type=sst, user=user)
+
+        out = self.call_command(f"--source={source_name}")
+
+        self.assertFalse(out)
+
+        self.assertTrue(
+            Incident.objects.exclude(id__in=previous_incidents_pks).filter(source__name=source_name).exists()
+        )
+
+    def test_create_fake_incident_will_raise_error_for_non_existent_source(self):
+        with self.assertRaises(CommandError):
+            self.call_command("--source=nonexistent")
+
     def test_create_fake_incident_will_create_single_fake_incident_with_set_level(self):
         previous_incidents_pks = [incident.id for incident in Incident.objects.all()]
         level = 2