Fix rule import for Project/Service

- Add test cases for each scenario
- Refactor settings for add_user_permissions
- Ensure rule is registered with correct ContentType

promgen/tests/__init__.py

@@ -3,7 +3,10 @@
 
 import json
 import os
+
 import yaml
+
+from django.contrib.auth.models import Permission
 from django.test import TestCase
 
 
@@ -36,3 +39,21 @@ def factory(self, klass, name):
                 name=name,
                 service=service,
             )
+
+    def assertRoute(self, response, view, status=200):
+        self.assertEqual(response.status_code, status)
+        self.assertEqual(response.resolver_match.func.__name__, view.as_view().__name__)
+
+    def assertCount(self, model, count, msg=None):
+        self.assertEqual(model.objects.count(), count, msg)
+
+    def add_user_permissions(self, *args, user=None):
+        codenames = [p.split(".")[1] for p in args]
+        permissions = Permission.objects.filter(
+            codename__in=codenames, content_type__app_label="promgen"
+        )
+
+        if user is None:
+            user = self.user
+
+        user.user_permissions.add(*[p for p in permissions])

promgen/tests/test_rules.py

@@ -4,10 +4,10 @@
 from unittest import mock
 
 import promgen.templatetags.promgen as macro
-from promgen import models, prometheus
+from promgen import models, prometheus, views
 from promgen.tests import PromgenTest
 
-from django.contrib.auth.models import Permission, User
+from django.contrib.auth.models import User
 from django.core.exceptions import ValidationError
 from django.test import override_settings
 from django.urls import reverse
@@ -63,20 +63,57 @@ def test_copy(self, mock_post):
         self.assertEqual(models.RuleLabel.objects.count(), 3, 'Copied rule has exiting labels + service label')
         self.assertEqual(models.RuleAnnotation.objects.count(), 2)
 
-    @mock.patch('django.dispatch.dispatcher.Signal.send')
+    @override_settings(PROMGEN=TEST_SETTINGS)
+    @mock.patch("django.dispatch.dispatcher.Signal.send")
     def test_import_v2(self, mock_post):
-        self.user.user_permissions.add(
-            Permission.objects.get(codename='change_rule'),
-            Permission.objects.get(codename='change_site', content_type__app_label='promgen'),
+        self.add_user_permissions("promgen.change_rule", "promgen.change_site")
+        response = self.client.post(
+            reverse("rule-import"),
+            {"rules": PromgenTest.data("examples", "import.rule.yml")},
+            follow=True,
         )
-        self.client.post(reverse('rule-import'), {
-            'rules': PromgenTest.data('examples', 'import.rule.yml')
-        })
 
         # Includes count of our setUp rule + imported rules
-        self.assertEqual(models.Rule.objects.count(), 3, 'Missing Rule')
-        self.assertEqual(models.RuleLabel.objects.count(), 4, 'Missing labels')
-        self.assertEqual(models.RuleAnnotation.objects.count(), 9, 'Missing annotations')
+        self.assertRoute(response, views.RuleImport, status=200)
+        self.assertCount(models.Rule, 3, "Missing Rule")
+        self.assertCount(models.RuleLabel, 4, "Missing labels")
+        self.assertCount(models.RuleAnnotation, 9, "Missing annotations")
+
+    @override_settings(PROMGEN=TEST_SETTINGS)
+    @mock.patch("django.dispatch.dispatcher.Signal.send")
+    def test_import_project_rule(self, mock_post):
+        self.add_user_permissions("promgen.add_rule", "promgen.change_project")
+        project = models.Project.objects.create(
+            name="Project 1", service=self.service, shard=self.shard
+        )
+        response = self.client.post(
+            reverse(
+                "rule-new", kwargs={"content_type": "project", "object_id": project.id}
+            ),
+            {"rules": PromgenTest.data("examples", "import.rule.yml")},
+            follow=True,
+        )
+        self.assertRoute(response, views.ProjectDetail, status=200)
+        self.assertCount(models.Rule, 3, "Missing Rule")
+        self.assertCount(models.RuleLabel, 4, "Missing labels")
+        self.assertCount(models.RuleAnnotation, 9, "Missing annotations")
+
+    @override_settings(PROMGEN=TEST_SETTINGS)
+    @mock.patch("django.dispatch.dispatcher.Signal.send")
+    def test_import_service_rule(self, mock_post):
+        self.add_user_permissions("promgen.add_rule", "promgen.change_service")
+        response = self.client.post(
+            reverse(
+                "rule-new",
+                kwargs={"content_type": "service", "object_id": self.service.id},
+            ),
+            {"rules": PromgenTest.data("examples", "import.rule.yml")},
+            follow=True,
+        )
+        self.assertRoute(response, views.ServiceDetail, status=200)
+        self.assertCount(models.Rule, 3, "Missing Rule")
+        self.assertCount(models.RuleLabel, 4, "Missing labels")
+        self.assertCount(models.RuleAnnotation, 9, "Missing annotations")
 
     @mock.patch('django.dispatch.dispatcher.Signal.send')
     def test_missing_permission(self, mock_post):

promgen/views.py

@@ -822,7 +822,7 @@ def post(self, request, *args, **kwargs):
         return self.form_valid(form)
 
 
-class RuleRegister(PromgenPermissionMixin, FormView, ServiceMixin):
+class RuleRegister(PromgenPermissionMixin, FormView):
     model = models.Rule
     template_name = 'promgen/rule_register.html'
     form_class = forms.RuleForm
@@ -857,12 +857,14 @@ def post(self, request, content_type, object_id):
             return self.form_invalid(form)
 
         importform = forms.ImportRuleForm(request.POST)
-        service = get_object_or_404(models.Service, id=self.kwargs['pk'])
+        ct = ContentType.objects.get_by_natural_key('promgen', content_type).model_class()
+        obj = ct.objects.get(pk=object_id)
+
         if importform.is_valid():
             data = importform.clean()
-            counters = prometheus.import_rules_v2(data['rules'], service)
+            counters = prometheus.import_rules_v2(data['rules'], obj)
             messages.info(request, 'Imported %s' % counters)
-            return HttpResponseRedirect(service.get_absolute_url())
+            return HttpResponseRedirect(obj.get_absolute_url())
 
         return self.form_invalid(form)