feat: endpoints with authorization, missing tests

Author: mjsifri

Diff for: authentication/decorators.py

@@ -7,14 +7,37 @@
 
 def cognito_authenticated(func):
     @wraps(func)
-    @api_view(["GET", "POST", "PUT", "DELETE"])
-    @authentication_classes([CognitoAuthentication])
-    def wrapper(request, *args, **kwargs):
-        if not request.user or not hasattr(request.user, "get"):
+    def wrapper(self, request, *args, **kwargs):
+        auth = CognitoAuthentication()
+        user, auth_error = auth.authenticate(request)
+        if auth_error or not user:
             return Response(
                 {"error": "Authentication credentials were not provided or are invalid."},
                 status=status.HTTP_401_UNAUTHORIZED,
             )
-        return func(request, *args, **kwargs)
+        request.user = user
+        return func(self, request, *args, **kwargs)
 
     return wrapper
+
+
+# from functools import wraps
+# from rest_framework.response import Response
+# from rest_framework import status
+# from rest_framework.decorators import api_view, authentication_classes
+# from .services.cognito_authentication import CognitoAuthentication
+
+
+# def cognito_authenticated(func):
+#     @wraps(func)
+#     @api_view(["GET", "POST", "PUT", "DELETE"])
+#     @authentication_classes([CognitoAuthentication])
+#     def wrapper(request, *args, **kwargs):
+#         if not request.user or not hasattr(request.user, "get"):
+#             return Response(
+#                 {"error": "Authentication credentials were not provided or are invalid."},
+#                 status=status.HTTP_401_UNAUTHORIZED,
+#             )
+#         return func(request, *args, **kwargs)
+
+#     return wrapper

Diff for: budget/admin.py

@@ -2,7 +2,8 @@
 from django.contrib import admin
 from .models import Budget
 
+
 @admin.register(Budget)
 class BudgetAdmin(admin.ModelAdmin):
-    list_display = ('user', 'amount', 'created_at', 'updated_at')
-    search_fields = ('user__username', 'amount')
+    list_display = ("username", "amount", "created_at", "updated_at")
+    search_fields = ("username", "amount")

Diff for: budget/migrations/0001_initial.py

@@ -1,27 +1,24 @@
-# Generated by Django 5.0.6 on 2024-06-10 00:41
+# Generated by Django 5.0.6 on 2024-06-10 15:06
 
-import django.db.models.deletion
-from django.conf import settings
+import uuid
 from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
 
     initial = True
 
-    dependencies = [
-        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
-    ]
+    dependencies = []
 
     operations = [
         migrations.CreateModel(
             name="Budget",
             fields=[
-                ("id", models.AutoField(primary_key=True, serialize=False)),
+                ("id", models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)),
+                ("username", models.UUIDField()),
                 ("amount", models.IntegerField()),
                 ("created_at", models.DateTimeField(auto_now_add=True)),
                 ("updated_at", models.DateTimeField(auto_now=True)),
-                ("user", models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
             ],
         ),
     ]

Diff for: budget/models.py

@@ -1,11 +1,12 @@
 # budget/models.py
 from django.db import models
 from django.conf import settings
+import uuid
 
 
 class Budget(models.Model):
-    id = models.AutoField(primary_key=True)
-    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+    username = models.UUIDField()
     amount = models.IntegerField()
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)

Diff for: budget/serializers.py

@@ -0,0 +1,10 @@
+from rest_framework import serializers
+from .models import Budget
+
+
+class BudgetSerializer(serializers.Serializer):
+    username = serializers.UUIDField()
+    amount = serializers.IntegerField()
+
+    def create(self, validated_data):
+        return Budget.objects.create(**validated_data)

Diff for: budget/tests.py

@@ -1,45 +1,51 @@
-# budget/tests.py
-from django.test import TestCase, Client
-from django.urls import reverse
-from django.contrib.auth import get_user_model
+from django.test import TestCase
+from rest_framework.test import APIRequestFactory
+from .views import BudgetViewSet
 from .models import Budget
-import json
+from django.contrib.auth.models import User
+from unittest.mock import patch
 
 
-class BudgetTests(TestCase):
+class BudgetViewSetTestCase(TestCase):
     def setUp(self):
-        self.client = Client()
-        self.user = get_user_model().objects.create_user(username="testuser", password="password")
+        self.factory = APIRequestFactory()
+        self.viewset = BudgetViewSet()
+        self.user = User.objects.create_user(username="testuser", password="testpassword")
         self.budget = Budget.objects.create(user=self.user, amount=100)
 
-    def test_get_user_budget(self):
-        url = reverse("get_user_budget")
-        response = self.client.get(url, {"user_id": self.user.id}, content_type="application/json")
+    @patch("budget.views.get_user_id_from_token")
+    def test_list(self, mock_get_user_id_from_token):
+        mock_get_user_id_from_token.return_value = self.user.id
+
+        request = self.factory.get("/")
+        response = self.viewset.list(request)
         self.assertEqual(response.status_code, 200)
-        self.assertEqual(response.json(), {"user_id": self.user.id, "amount": self.budget.amount})
-
-    def test_set_user_budget(self):
-        url = reverse("set_user_budget")
-        user = get_user_model().objects.create_user(username="testuser2", password="password")
-        response = self.client.post(
-            url, json.dumps({"user_id": user.id, "amount": 200}), content_type="application/json"
-        )
+
+    @patch("budget.views.get_user_id_from_token")
+    def test_create(self, mock_get_user_id_from_token):
+        mock_get_user_id_from_token.return_value = self.user.id
+
+        request_data = {"amount": 200}
+        request = self.factory.post("/", request_data, format="json")
+        response = self.viewset.create(request)
         self.assertEqual(response.status_code, 201)
-        self.assertEqual(response.json(), {"message": "Budget created successfully."})
-        self.assertEqual(Budget.objects.get(user=user).amount, 200)
+        self.assertEqual(Budget.objects.count(), 2)
 
-    def test_delete_user_budget(self):
-        url = reverse("delete_user_budget")
-        response = self.client.delete(url, json.dumps({"user_id": self.user.id}), content_type="application/json")
-        self.assertEqual(response.status_code, 200)
-        self.assertEqual(response.json(), {"message": "Budget deleted successfully", "user_id": self.user.id})
-        self.assertFalse(Budget.objects.filter(user=self.user).exists())
-
-    def test_update_user_budget(self):
-        url = reverse("update_user_budget")
-        response = self.client.patch(
-            url, json.dumps({"user_id": self.user.id, "amount": 300}), content_type="application/json"
-        )
+    @patch("budget.views.get_user_id_from_token")
+    def test_destroy(self, mock_get_user_id_from_token):
+        mock_get_user_id_from_token.return_value = self.user.id
+
+        request = self.factory.delete("/")
+        response = self.viewset.destroy(request, pk=self.budget.id)
+        self.assertEqual(response.status_code, 204)
+        self.assertEqual(Budget.objects.count(), 0)
+
+    @patch("budget.views.get_user_id_from_token")
+    def test_partial_update(self, mock_get_user_id_from_token):
+        mock_get_user_id_from_token.return_value = self.user.id
+
+        request_data = {"amount": 150}
+        request = self.factory.patch("/", request_data, format="json")
+        response = self.viewset.partial_update(request, pk=self.budget.id)
         self.assertEqual(response.status_code, 200)
-        self.assertEqual(response.json(), {"user_id": self.user.id, "budget": 300})
-        self.assertEqual(Budget.objects.get(user=self.user).amount, 300)
+        self.assertEqual(Budget.objects.get(pk=self.budget.id).amount, 150)

Diff for: budget/urls.py

@@ -1,10 +1,11 @@
 # budget/urls.py
 from django.urls import path
-from .views import get_user_budget, set_user_budget, delete_user_budget, update_user_budget
+from .views import BudgetViewSet
 
 urlpatterns = [
-    path("get", get_user_budget, name="get_user_budget"),
-    path("set", set_user_budget, name="set_user_budget"),
-    path("delete", delete_user_budget, name="delete_user_budget"),
-    path("update", update_user_budget, name="update_user_budget"),
+    path(
+        "",
+        BudgetViewSet.as_view({"get": "list", "post": "create", "delete": "destroy", "put": "partial_update"}),
+        name="budget",
+    ),
 ]

Diff for: budget/views.py

@@ -1,85 +1,82 @@
-# budget/views.py
-from django.shortcuts import render
-from django.http import JsonResponse
+from rest_framework import viewsets, status
+from rest_framework.response import Response
 from .models import Budget
-from django.views.decorators.csrf import csrf_exempt
-import json
+from .serializers import BudgetSerializer
+from authentication.decorators import cognito_authenticated
+import jwt
 
 
-@csrf_exempt
-def get_user_budget(request):
-    if request.method == "GET":
+class BudgetViewSet(viewsets.ViewSet):
+    def get_user_id_from_token(self, request):
         try:
-            user_id = request.GET.get("user_id")
+            authorization_header = request.headers.get("Authorization")
+            if not authorization_header:
+                raise Exception("Authorization header not found")
+
+            token = authorization_header.split()[1]
+            decoded_token = jwt.decode(token, options={"verify_signature": False})
+            user_id = decoded_token.get("username")
             if not user_id:
-                return JsonResponse({"error": "User ID is required in the query parameters."}, status=400)
+                raise Exception("User ID not found in token")
+            return user_id
+        except jwt.DecodeError:
+            raise Exception("Invalid token")
+        except jwt.ExpiredSignatureError:
+            raise Exception("Expired token")
+        except Exception as e:
+            raise Exception(f"Error decoding token: {e}")
 
-            budget = Budget.objects.get(user_id=user_id)
-            return JsonResponse({"user_id": int(user_id), "amount": budget.amount})
+    @cognito_authenticated
+    def list(self, request):
+        try:
+            username = self.get_user_id_from_token(request)
+            budget = Budget.objects.get(username=username)
+            serializer = BudgetSerializer(budget)
+            return Response(serializer.data)
         except Budget.DoesNotExist:
-            return JsonResponse({"error": "Budget not found"}, status=404)
-    else:
-        return JsonResponse({"error": "Invalid request method."}, status=405)
-
+            return Response({"error": "Budget not found"}, status=status.HTTP_404_NOT_FOUND)
+        except Exception as e:
+            return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
-@csrf_exempt
-def set_user_budget(request):
-    if request.method == "POST":
+    @cognito_authenticated
+    def create(self, request):
         try:
-            data = json.loads(request.body)
-            user_id = data.get("user_id")
-            amount = data.get("amount")
-
-            if not user_id or not amount:
-                return JsonResponse({"error": "User ID and amount are required."}, status=400)
-
-            budget = Budget.objects.create(user_id=user_id, amount=amount)
-            budget.save()
-
-            return JsonResponse({"message": "Budget created successfully."}, status=201)
-        except json.JSONDecodeError:
-            return JsonResponse({"error": "Invalid JSON."}, status=400)
+            user_id = self.get_user_id_from_token(request)
+            data = request.data.copy()
+            data["username"] = user_id
+
+            serializer = BudgetSerializer(data=data)
+            print(serializer)
+            if serializer.is_valid():
+                serializer.save()
+                return Response(serializer.data, status=status.HTTP_201_CREATED)
+            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
         except Exception as e:
-            return JsonResponse({"error": str(e)}, status=500)
-    else:
-        return JsonResponse({"error": "Invalid request method."}, status=405)
+            return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
-
-@csrf_exempt
-def delete_user_budget(request):
-    if request.method == "DELETE":
+    @cognito_authenticated
+    def destroy(self, request):
         try:
-            data = json.loads(request.body)
-            user_id = data.get("user_id")
-
-            budget = Budget.objects.get(user_id=user_id)
+            username = self.get_user_id_from_token(request)
+            budget = Budget.objects.get(username=username)
             budget.delete()
-
-            return JsonResponse({"message": "Budget deleted successfully", "user_id": user_id})
+            return Response(status=status.HTTP_204_NO_CONTENT)
         except Budget.DoesNotExist:
-            return JsonResponse({"error": "Budget not found"}, status=404)
-        except json.JSONDecodeError:
-            return JsonResponse({"error": "Invalid JSON data"}, status=400)
-    else:
-        return JsonResponse({"error": "Invalid request method."}, status=405)
-
+            return Response({"error": "Budget not found"}, status=status.HTTP_404_NOT_FOUND)
+        except Exception as e:
+            return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
-@csrf_exempt
-def update_user_budget(request):
-    if request.method == "PATCH":
+    @cognito_authenticated
+    def partial_update(self, request):
         try:
-            data = json.loads(request.body)
-            user_id = data.get("user_id")
-            amount = data.get("amount")
-
-            budget = Budget.objects.get(user_id=user_id)
-            budget.amount = amount
-            budget.save()
-
-            return JsonResponse({"user_id": user_id, "budget": budget.amount})
+            username = self.get_user_id_from_token(request)
+            budget = Budget.objects.get(username=username)
+            serializer = BudgetSerializer(budget, data=request.data, partial=True)
+            if serializer.is_valid():
+                serializer.save()
+                return Response(serializer.data)
+            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
         except Budget.DoesNotExist:
-            return JsonResponse({"error": "Budget not found"}, status=404)
-        except json.JSONDecodeError:
-            return JsonResponse({"error": "Invalid JSON data"}, status=400)
-    else:
-        return JsonResponse({"error": "Invalid request method."}, status=405)
+            return Response({"error": "Budget not found"}, status=status.HTTP_404_NOT_FOUND)
+        except Exception as e:
+            return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)