wip

staff/models.py

@@ -11,7 +11,7 @@ def __str__(self) -> str:
         return f"{self.employee_no} - {self.first_name} {self.last_name}"
 
 
-class StaffForecast(models.QuerySet):
+class StaffForecastQuerySet(models.QuerySet):
     pass
 
 
@@ -24,7 +24,7 @@ class Meta:
             )
         ]
 
-    objects = StaffForecast.as_manager()
+    objects = StaffForecastQuerySet.as_manager()
 
     staff = models.ForeignKey(Staff, models.PROTECT, related_name="forecast")
     year = models.ForeignKey("core.FinancialYear", models.PROTECT)

staff/tests/test_views.py

@@ -0,0 +1,30 @@
+import pytest
+
+from django.contrib.auth import get_user_model
+
+
+User = get_user_model()
+
+
+@pytest.fixture
+def user(db, client):
+    user = User.objects.create_user(
+        username="staff.test",
+        email="[email protected]",
+        password="password",
+    )
+    user.save()
+    client.force_login(user)
+    return user
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "/staff/edit-payroll/",
+        "/staff/debug/",
+    ],
+)
+def test_only_superuser_can_access(client, user, url):
+    r = client.get(url)
+    assert r.status_code == 403

staff/views.py

@@ -1,6 +1,8 @@
+from functools import wraps
 from django.http import HttpResponse, HttpRequest
 from django.template.response import TemplateResponse
 from django.contrib.auth.decorators import user_passes_test
+from django.core.exceptions import PermissionDenied
 
 from core.models import FinancialYear
 from costcentre.models import CostCentre
@@ -10,17 +12,23 @@
 
 
 # TODO: Remove once no longer needed.
-def _user_is_superuser(user):
-    return user.is_superuser
+def superuser_view(view_func):
+    @wraps(view_func)
+    def wrapper(request, *args, **kwargs):
+        if not request.user.is_superuser:
+            raise PermissionDenied
+        return view_func(*args, **kwargs)
 
+    return wrapper
 
-@user_passes_test(_user_is_superuser)
+
+@superuser_view
 def edit_payroll_page(request: HttpRequest) -> HttpResponse:
     context = {}
     return TemplateResponse(request, "staff/page/edit_payroll.html", context)
 
 
-@user_passes_test(_user_is_superuser)
+@superuser_view
 def staff_debug_page(request: HttpRequest) -> HttpResponse:
     if request.GET.get("cost_centre"):
         cost_centre = CostCentre.objects.get(pk=request.GET.get("cost_centre"))