[FIX] Store the correct rule ID in rule annotation #556

promgen/models.py

@@ -485,6 +485,23 @@ def copy_to(self, content_type, object_id):
 
         return self
 
+    # Custom logic before saving Rule to control the value of the annotation "rule":
+    # Format: annotations["rule"] = {domain}/rule/{id}
+    def save(self, *args, **kwargs):
+        with transaction.atomic():
+            if self.pk:
+                # When updating rule, we already have the primary key.
+                # Just set annotations["rule"] before saving to database.
+                self.annotations["rule"] = resolve_domain("rule-detail", pk=self.pk)
+                super().save(*args, **kwargs)
+            else:
+                # When creating a new rule, the primary key is typically not available until the
+                # instance is saved to the database.
+                # Therefore, we save it first then set annotations["rule"] and save again.
+                super().save(*args, **kwargs)
+                self.annotations["rule"] = resolve_domain("rule-detail", pk=self.pk)
+                super().save(update_fields=["annotations"])
+
 
 class AlertLabel(models.Model):
     alert = models.ForeignKey("Alert", on_delete=models.CASCADE)

promgen/tests/test_models.py

@@ -1,10 +1,12 @@
 # Copyright (c) 2020 LINE Corporation
 # These sources are released under the terms of the MIT license: see LICENSE
+from unittest import mock
 
-
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 
 from promgen import models, validators
+from promgen.shortcuts import resolve_domain
 from promgen.tests import PromgenTest
 
 
@@ -34,3 +36,36 @@ def test_validators(self):
             validators.metricname(
                 "[[this.$el.ownerDocument.defaultView.alert(1337)]]",
             )
+
+    @mock.patch("django.dispatch.dispatcher.Signal.send")
+    def test_rule_annotation(self, mock_post):
+        # Check if annotation["rule"] is automatically set to be {domain}/rule/{id} when creating a
+        # new rule
+        rule = models.Rule(
+            name="example-rule",
+            content_type=ContentType.objects.get_for_model(models.Site),
+            object_id=1,
+            clause="up==1",
+            duration="1s"
+        )
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] is automatically set to be {domain}/rule/{id} when updating an
+        # existed rule
+        rule.name = "another-example-rule"
+        rule.annotations["rule"] = "another-annotation-value"
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] is still set to be {domain}/rule/{id} when trying to remove
+        # annotation["rule"]
+        rule.annotations["rule"] = None
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] of new rule is automatically set to be {domain}/rule/{id}
+        # when cloning an existed rule
+        new_rule = rule.copy_to(content_type="service", object_id=2)
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+        self.assertEqual(resolve_domain("rule-detail", new_rule.pk), new_rule.annotations["rule"])