diff --git a/promgen/models.py b/promgen/models.py index 27330f29..d99467a8 100644 --- a/promgen/models.py +++ b/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) diff --git a/promgen/tests/test_models.py b/promgen/tests/test_models.py index 373146d6..33d42362 100644 --- a/promgen/tests/test_models.py +++ b/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"])