refactor(business auth): Consolidate JWT token utilities

- Simplify views using `CreateAPIView` and remove redundant mixins
- Add `@transaction.atomic` for atomic database operations during company creation
- Extract JWT token generation into reusable `generate_company_tokens()` utility
- Extract JWT token version bump into reusable `bump_company_token_version()` utility
- Validate UUID format rigorously in token refresh flow
- Add documentation in views
- Remove redundant imports

Author: RandomProgramm3r

promo_code/business/serializers.py

@@ -1,8 +1,8 @@
 import uuid
 
 import django.contrib.auth.password_validation
-import django.core.exceptions
 import django.core.validators
+import django.db.transaction
 import pycountry
 import rest_framework.exceptions
 import rest_framework.serializers
@@ -11,7 +11,10 @@
 import rest_framework_simplejwt.tokens
 
 import business.constants
+import business.models
 import business.models as business_models
+import business.utils.auth
+import business.utils.tokens
 import business.validators
 
 
@@ -21,9 +24,9 @@ class CompanySignUpSerializer(rest_framework.serializers.ModelSerializer):
         write_only=True,
         required=True,
         validators=[django.contrib.auth.password_validation.validate_password],
+        style={'input_type': 'password'},
         min_length=business.constants.COMPANY_PASSWORD_MIN_LENGTH,
         max_length=business.constants.COMPANY_PASSWORD_MAX_LENGTH,
-        style={'input_type': 'password'},
     )
     name = rest_framework.serializers.CharField(
         required=True,
@@ -44,30 +47,18 @@ class CompanySignUpSerializer(rest_framework.serializers.ModelSerializer):
 
     class Meta:
         model = business_models.Company
-        fields = (
-            'id',
-            'name',
-            'email',
-            'password',
-        )
+        fields = ('id', 'name', 'email', 'password')
 
+    @django.db.transaction.atomic
     def create(self, validated_data):
-        try:
-            company = business_models.Company.objects.create_company(
-                email=validated_data['email'],
-                name=validated_data['name'],
-                password=validated_data['password'],
-            )
-            company.token_version += 1
-            company.save()
-            return company
-        except django.core.exceptions.ValidationError as e:
-            raise rest_framework.serializers.ValidationError(e.messages)
+        company = business_models.Company.objects.create_company(
+            **validated_data,
+        )
 
+        return business.utils.auth.bump_company_token_version(company)
 
-class CompanySignInSerializer(
-    rest_framework.serializers.Serializer,
-):
+
+class CompanySignInSerializer(rest_framework.serializers.Serializer):
     email = rest_framework.serializers.EmailField(required=True)
     password = rest_framework.serializers.CharField(
         required=True,
@@ -80,16 +71,15 @@ def validate(self, attrs):
         password = attrs.get('password')
 
         if not email or not password:
-            raise rest_framework.exceptions.ValidationError(
-                {'detail': 'Both email and password are required'},
-                code='required',
+            raise rest_framework.serializers.ValidationError(
+                'Both email and password are required.',
             )
 
         try:
             company = business_models.Company.objects.get(email=email)
         except business_models.Company.DoesNotExist:
             raise rest_framework.serializers.ValidationError(
-                'Invalid credentials',
+                'Invalid credentials.',
             )
 
         if not company.is_active or not company.check_password(password):
@@ -98,53 +88,55 @@ def validate(self, attrs):
                 code='authentication_failed',
             )
 
+        attrs['company'] = company
         return attrs
 
 
 class CompanyTokenRefreshSerializer(
     rest_framework_simplejwt.serializers.TokenRefreshSerializer,
 ):
     def validate(self, attrs):
+        attrs = super().validate(attrs)
         refresh = rest_framework_simplejwt.tokens.RefreshToken(
             attrs['refresh'],
         )
-        user_type = refresh.payload.get('user_type', 'user')
+        company = self.get_active_company_from_token(refresh)
+
+        company = business.utils.auth.bump_company_token_version(company)
 
-        if user_type != 'company':
+        return business.utils.tokens.generate_company_tokens(company)
+
+    def get_active_company_from_token(self, token):
+        if token.payload.get('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:
+        company_id = token.payload.get('company_id')
+        try:
+            company_uuid = uuid.UUID(company_id)
+        except (TypeError, ValueError):
             raise rest_framework_simplejwt.exceptions.InvalidToken(
-                'Company ID missing in token',
+                'Invalid or missing company_id in token',
             )
 
         try:
-            company = business_models.Company.objects.get(
-                id=uuid.UUID(company_id),
+            company = business.models.Company.objects.get(
+                id=company_uuid,
+                is_active=True,
             )
-        except business_models.Company.DoesNotExist:
+        except business.models.Company.DoesNotExist:
             raise rest_framework_simplejwt.exceptions.InvalidToken(
-                'Company not found',
+                'Company not found or inactive',
             )
 
-        token_version = refresh.payload.get('token_version', 0)
+        token_version = token.payload.get('token_version', 0)
         if company.token_version != token_version:
             raise rest_framework_simplejwt.exceptions.InvalidToken(
                 'Token is blacklisted',
             )
 
-        new_refresh = rest_framework_simplejwt.tokens.RefreshToken()
-        new_refresh['user_type'] = 'company'
-        new_refresh['company_id'] = str(company.id)
-        new_refresh['token_version'] = company.token_version
-
-        return {
-            'access': str(new_refresh.access_token),
-            'refresh': str(new_refresh),
-        }
+        return company
 
 
 class TargetSerializer(rest_framework.serializers.Serializer):

promo_code/business/utils/auth.py

@@ -0,0 +1,13 @@
+import business.models
+
+
+def bump_company_token_version(company):
+    """
+    Increment token_version, save it, and return the fresh instance.
+    """
+    company = business.models.Company.objects.select_for_update().get(
+        id=company.id,
+    )
+    company.token_version += 1
+    company.save(update_fields=['token_version'])
+    return company

promo_code/business/utils/tokens.py

@@ -0,0 +1,20 @@
+import rest_framework_simplejwt.tokens
+
+
+def generate_company_tokens(company):
+    """
+    Generate JWT tokens for a company.
+    """
+    refresh = rest_framework_simplejwt.tokens.RefreshToken()
+    refresh['user_type'] = 'company'
+    refresh['company_id'] = str(company.id)
+    refresh['token_version'] = company.token_version
+
+    access = refresh.access_token
+    access['user_type'] = 'company'
+    access['company_id'] = str(company.id)
+
+    return {
+        'access': str(access),
+        'refresh': str(refresh),
+    }

promo_code/business/views.py

@@ -4,111 +4,61 @@
 import rest_framework.generics
 import rest_framework.permissions
 import rest_framework.response
-import rest_framework.serializers
 import rest_framework.status
-import rest_framework_simplejwt.exceptions
-import rest_framework_simplejwt.tokens
 import rest_framework_simplejwt.views
 
 import business.models
 import business.pagination
 import business.permissions
 import business.serializers
-import core.views
+import business.utils.auth
+import business.utils.tokens
 
 
-class CompanySignUpView(
-    core.views.BaseCustomResponseMixin,
-    rest_framework.generics.CreateAPIView,
-):
-    serializer_class = business.serializers.CompanySignUpSerializer
-
-    def post(self, request):
-        try:
-            serializer = business.serializers.CompanySignUpSerializer(
-                data=request.data,
-            )
-            serializer.is_valid(raise_exception=True)
-        except (
-            rest_framework.serializers.ValidationError,
-            rest_framework_simplejwt.exceptions.TokenError,
-        ) as e:
-            if isinstance(e, rest_framework.serializers.ValidationError):
-                return self.handle_validation_error()
+class CompanySignUpView(rest_framework.generics.CreateAPIView):
+    """
+    Company registration endpoint that returns JWT tokens.
+    """
 
-            raise rest_framework_simplejwt.exceptions.InvalidToken(str(e))
+    serializer_class = business.serializers.CompanySignUpSerializer
 
+    def create(self, request, *args, **kwargs):
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
         company = serializer.save()
 
-        refresh = rest_framework_simplejwt.tokens.RefreshToken()
-        refresh['user_type'] = 'company'
-        refresh['company_id'] = str(company.id)
-        refresh['token_version'] = company.token_version
-
-        access_token = refresh.access_token
-        access_token['user_type'] = 'company'
-        access_token['company_id'] = str(company.id)
-        refresh['token_version'] = company.token_version
-
-        response_data = {
-            'access': str(access_token),
-            'refresh': str(refresh),
-        }
-
         return rest_framework.response.Response(
-            response_data,
+            business.utils.tokens.generate_company_tokens(company),
             status=rest_framework.status.HTTP_200_OK,
         )
 
 
-class CompanySignInView(
-    core.views.BaseCustomResponseMixin,
-    rest_framework_simplejwt.views.TokenObtainPairView,
-):
-    serializer_class = business.serializers.CompanySignInSerializer
-
-    def post(self, request):
-        try:
-            serializer = business.serializers.CompanySignInSerializer(
-                data=request.data,
-            )
-            serializer.is_valid(raise_exception=True)
-        except (
-            rest_framework.serializers.ValidationError,
-            rest_framework_simplejwt.exceptions.TokenError,
-        ) as e:
-            if isinstance(e, rest_framework.serializers.ValidationError):
-                return self.handle_validation_error()
-
-            raise rest_framework_simplejwt.exceptions.InvalidToken(str(e))
-
-        company = business.models.Company.objects.get(
-            email=serializer.validated_data['email'],
-        )
-        company.token_version += 1
-        company.save()
+class CompanySignInView(rest_framework.generics.GenericAPIView):
+    """
+    Company authentication endpoint that issues new JWT tokens
+    and bumps token_version.
+    """
 
-        refresh = rest_framework_simplejwt.tokens.RefreshToken()
-        refresh['user_type'] = 'company'
-        refresh['company_id'] = str(company.id)
-        refresh['token_version'] = company.token_version
+    serializer_class = business.serializers.CompanySignInSerializer
 
-        access_token = refresh.access_token
-        access_token['user_type'] = 'company'
-        access_token['company_id'] = str(company.id)
+    def post(self, request, *args, **kwargs):
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
 
-        response_data = {
-            'access': str(access_token),
-            'refresh': str(refresh),
-        }
+        company = serializer.validated_data['company']
+        company = business.utils.auth.bump_company_token_version(company)
 
         return rest_framework.response.Response(
-            response_data,
+            business.utils.tokens.generate_company_tokens(company),
             status=rest_framework.status.HTTP_200_OK,
         )
 
 
 class CompanyTokenRefreshView(rest_framework_simplejwt.views.TokenRefreshView):
+    """
+    Refresh endpoint for company tokens only.
+    """
+
     serializer_class = business.serializers.CompanyTokenRefreshSerializer