feat: Add the ability to update the token for the company.

Author: RandomProgramm3r

promo_code/business/serializers.py

@@ -6,6 +6,10 @@
 import rest_framework.exceptions
 import rest_framework.serializers
 import rest_framework.status
+import rest_framework_simplejwt.exceptions
+import rest_framework_simplejwt.serializers
+import rest_framework_simplejwt.tokens
+import rest_framework_simplejwt.views
 
 
 class CompanySignUpSerializer(rest_framework.serializers.ModelSerializer):
@@ -90,3 +94,39 @@ def validate(self, attrs):
             )
 
         return attrs
+
+
+class CompanyTokenRefreshSerializer(
+    rest_framework_simplejwt.serializers.TokenRefreshSerializer,
+):
+    def validate(self, attrs):
+        refresh = rest_framework_simplejwt.tokens.RefreshToken(
+            attrs['refresh'],
+        )
+        user_type = refresh.payload.get('user_type', 'user')
+
+        if user_type != 'company':
+            raise rest_framework_simplejwt.exceptions.InvalidToken(
+                'This refresh endpoint is for company tokens only',
+            )
+
+        company_id = refresh.payload.get('company_id')
+        if not company_id:
+            raise rest_framework_simplejwt.exceptions.InvalidToken(
+                'Company ID missing in token',
+            )
+
+        try:
+            company = business_models.Company.objects.get(id=company_id)
+        except business_models.Company.DoesNotExist:
+            raise rest_framework_simplejwt.exceptions.InvalidToken(
+                'Company not found',
+            )
+
+        token_version = refresh.payload.get('token_version', 0)
+        if company.token_version != token_version:
+            raise rest_framework_simplejwt.exceptions.InvalidToken(
+                'Token is blacklisted',
+            )
+
+        return super().validate(attrs)

promo_code/business/tests/auth/base.py

@@ -10,9 +10,12 @@ class BaseBusinessAuthTestCase(rest_framework.test.APITestCase):
     def setUpTestData(cls):
         super().setUpTestData()
         cls.client = rest_framework.test.APIClient()
+        cls.company_refresh_url = django.urls.reverse(
+            'api-business:company-token-refresh',
+        )
+        cls.protected_url = django.urls.reverse('api-core:protected')
         cls.signup_url = django.urls.reverse('api-business:company-sign-up')
         cls.signin_url = django.urls.reverse('api-business:company-sign-in')
-        cls.protected_url = django.urls.reverse('api-core:protected')
         cls.valid_data = {
             'name': 'Digital Marketing Solutions Inc.',
             'email': 'testcompany@example.com',

promo_code/business/urls.py

@@ -15,4 +15,9 @@
         business.views.CompanySignInView.as_view(),
         name='company-sign-in',
     ),
+    django.urls.path(
+        'token/refresh',
+        business.views.CompanyTokenRefreshView.as_view(),
+        name='company-token-refresh',
+    ),
 ]

promo_code/business/views.py

@@ -97,3 +97,7 @@ def post(self, request):
             response_data,
             status=rest_framework.status.HTTP_200_OK,
         )
+
+
+class CompanyTokenRefreshView(rest_framework_simplejwt.views.TokenRefreshView):
+    serializer_class = business.serializers.CompanyTokenRefreshSerializer

promo_code/user/tests/auth/base.py

@@ -11,7 +11,7 @@ def setUpTestData(cls):
         super().setUpTestData()
         cls.client = rest_framework.test.APIClient()
         cls.protected_url = django.urls.reverse('api-core:protected')
-        cls.refresh_url = django.urls.reverse('api-user:token_refresh')
+        cls.refresh_url = django.urls.reverse('api-user:user-token-refresh')
         cls.signup_url = django.urls.reverse('api-user:sign-up')
         cls.signin_url = django.urls.reverse('api-user:sign-in')
 

promo_code/user/urls.py

@@ -20,6 +20,6 @@
     django.urls.path(
         'token/refresh/',
         rest_framework_simplejwt.views.TokenRefreshView.as_view(),
-        name='token_refresh',
+        name='user-token-refresh',
     ),
 ]