fix: Replace reference to User model with get_user_model()

dragonejt · Jun 23, 2024 · 07a31e1 · 07a31e1
1 parent 6870958
commit 07a31e1
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 17 deletions.
diff --git a/community/tests.py b/community/tests.py
@@ -1,6 +1,6 @@
 from rest_framework.test import APITestCase
 from rest_framework import status
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 from django.utils.crypto import get_random_string
 from rest_framework.authtoken.models import Token
 from community.views import Community, CommunitySerializer
@@ -12,7 +12,7 @@ class TestCommunityView(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/community"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)
@@ -29,11 +29,23 @@ def test_get(self) -> None:
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.json(), CommunitySerializer(self.community).data)
 
-    def test_head(self) -> None:
+    def test_head_ok(self) -> None:
         response = self.client.head(f"{self.url}?id={self.community.community_id}")
 
         self.assertEqual(response.status_code, status.HTTP_200_OK)
 
+    def test_head_forbidden(self) -> None:
+        user = get_user_model().objects.create(
+            username=get_random_string(10),
+            email=get_random_string(10),
+            password=get_random_string(10),
+        )
+        community = Community.objects.create(
+            platform=user, community_id=get_random_string(20)
+        )
+        response = self.client.head(f"{self.url}?id={community.community_id}")
+        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
     def test_post(self) -> None:
         community_id = get_random_string(20)
         response = self.client.post(

diff --git a/community/views.py b/community/views.py
@@ -35,13 +35,9 @@ def post(self, request: Request) -> Response:
         community = Community.objects.create(
             platform=request.user, community_id=request.data.get("communityID")
         )
-        community.save()
-        community_psycho_pass = CommunityPsychoPass.objects.create(community=community)
-        community_psycho_pass.save()
-        member_dominator = MemberDominator.objects.create(community=community)
-        member_dominator.save()
-        message_dominator = MessageDominator.objects.create(community=community)
-        message_dominator.save()
+        CommunityPsychoPass.objects.create(community=community)
+        MemberDominator.objects.create(community=community)
+        MessageDominator.objects.create(community=community)
 
         return Response(
             CommunitySerializer(community).data, status=status.HTTP_201_CREATED

diff --git a/dominator/tests.py b/dominator/tests.py
@@ -1,7 +1,7 @@
 import random
 from rest_framework.test import APITestCase
 from rest_framework import status
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 from django.utils.crypto import get_random_string
 from rest_framework.authtoken.models import Token
 from dominator.views import (
@@ -19,7 +19,7 @@ class TestMemberDominatorView(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/dominator/member"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)
@@ -84,7 +84,7 @@ class TestMessageDominatorView(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/dominator/message"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)

diff --git a/psychopass/tests.py b/psychopass/tests.py
@@ -1,6 +1,6 @@
 from random import random
 from rest_framework.test import APITestCase
-from django.contrib.auth.models import User
+from django.contrib.auth import get_user_model
 from django.utils.crypto import get_random_string
 from rest_framework import status
 from rest_framework.authtoken.models import Token
@@ -19,7 +19,7 @@ class TestUserPsychoPassView(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/psychopass/user"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)
@@ -62,7 +62,7 @@ class TestCommunityPsychoPassView(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/psychopass/community"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)
@@ -111,7 +111,7 @@ class TestIngestMessage(APITestCase):
 
     def setUp(self) -> None:
         self.url = "/psychopass/message"
-        self.user = User.objects.create_superuser(
+        self.user = get_user_model().objects.create_superuser(
             username=get_random_string(10), email=None, password=None
         )
         self.token = Token.objects.create(user=self.user)