Add Tests for PR#1101

VERSION

@@ -1 +1 @@
-25.1.10
+25.2.3

api/gen_ai/serializers.py

@@ -0,0 +1,5 @@
+from rest_framework import serializers
+
+
+class GenAIAuthSerializer(serializers.Serializer):
+    is_valid = serializers.BooleanField()

api/gen_ai/tests/test_gen_ai.py

@@ -0,0 +1,114 @@
+import hmac
+from hashlib import sha256
+from unittest.mock import patch
+
+from django.urls import reverse
+from rest_framework import status
+from rest_framework.test import APITestCase
+from shared.django_apps.core.tests.factories import OwnerFactory
+
+from codecov_auth.models import GithubAppInstallation
+
+PAYLOAD_SECRET = b"testixik8qdauiab1yiffydimvi72ekq"
+VIEW_URL = reverse("auth")
+
+
+def sign_payload(data: bytes, secret=PAYLOAD_SECRET):
+    signature = "sha256=" + hmac.new(secret, data, digestmod=sha256).hexdigest()
+    return signature, data
+
+
+class GenAIAuthViewTests(APITestCase):
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_missing_parameters(self, mock_config):
+        payload = b"{}"
+        sig, data = sign_payload(payload)
+        response = self.client.post(
+            VIEW_URL,
+            data=data,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
+        )
+        self.assertEqual(response.status_code, 400)
+        self.assertIn("Missing required parameters", response.data)
+
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_invalid_signature(self, mock_config):
+        # Correct payload
+        payload = b'{"external_owner_id":"owner1","repo_service_id":"101"}'
+        # Wrong signature based on a different payload
+        wrong_sig = "sha256=" + hmac.new(PAYLOAD_SECRET, b"{}", sha256).hexdigest()
+        response = self.client.post(
+            VIEW_URL,
+            data=payload,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=wrong_sig,
+        )
+        self.assertEqual(response.status_code, 403)
+
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_owner_not_found(self, mock_config):
+        payload = b'{"external_owner_id":"nonexistent_owner","repo_service_id":"101"}'
+        sig, data = sign_payload(payload)
+        response = self.client.post(
+            VIEW_URL,
+            data=data,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
+        )
+        self.assertEqual(response.status_code, 404)
+
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_no_installation(self, mock_config):
+        # Create a valid owner but no installation
+        OwnerFactory(service="github", service_id="owner1", username="test1")
+        payload = b'{"external_owner_id":"owner1","repo_service_id":"101"}'
+        sig, data = sign_payload(payload)
+        response = self.client.post(
+            VIEW_URL,
+            data=data,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.data, {"is_valid": False})
+
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_authorized(self, mock_config):
+        owner = OwnerFactory(service="github", service_id="owner2", username="test2")
+        GithubAppInstallation.objects.create(
+            installation_id=12345,
+            owner=owner,
+            name="ai-features",
+            repository_service_ids=["101", "202"],
+        )
+        payload = b'{"external_owner_id":"owner2","repo_service_id":"101"}'
+        sig, data = sign_payload(payload)
+        response = self.client.post(
+            VIEW_URL,
+            data=data,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
+        )
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(response.data, {"is_valid": True})
+
+    @patch("api.gen_ai.views.get_config", return_value=PAYLOAD_SECRET)
+    def test_unauthorized(self, mock_config):
+        owner = OwnerFactory(service="github", service_id="owner3", username="test3")
+        GithubAppInstallation.objects.create(
+            installation_id=2,
+            owner=owner,
+            name="ai-features",
+            repository_service_ids=["303", "404"],
+        )
+        payload = b'{"external_owner_id":"owner3","repo_service_id":"101"}'
+        sig, data = sign_payload(payload)
+        response = self.client.post(
+            VIEW_URL,
+            data=data,
+            content_type="application/json",
+            HTTP_HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.data, {"is_valid": False})

api/gen_ai/urls.py

@@ -0,0 +1,7 @@
+from django.urls import path
+
+from .views import GenAIAuthView
+
+urlpatterns = [
+    path("auth/", GenAIAuthView.as_view(), name="auth"),
+]

api/gen_ai/views.py

@@ -0,0 +1,61 @@
+import hmac
+import logging
+from hashlib import sha256
+
+from rest_framework.exceptions import NotFound, PermissionDenied
+from rest_framework.permissions import AllowAny
+from rest_framework.response import Response
+from rest_framework.views import APIView
+
+from api.gen_ai.serializers import GenAIAuthSerializer
+from codecov_auth.models import GithubAppInstallation, Owner
+from graphql_api.types.owner.owner import AI_FEATURES_GH_APP_ID
+from utils.config import get_config
+
+log = logging.getLogger(__name__)
+
+
+class GenAIAuthView(APIView):
+    permission_classes = [AllowAny]
+    serializer_class = GenAIAuthSerializer
+
+    def validate_signature(self, request):
+        key = get_config("gen_ai", "auth_secret")
+        if not key:
+            raise PermissionDenied("Invalid signature")
+
+        if isinstance(key, str):
+            key = key.encode("utf-8")
+        expected_sig = request.headers.get("HTTP-X-GEN-AI-AUTH-SIGNATURE")
+        computed_sig = (
+            "sha256=" + hmac.new(key, request.body, digestmod=sha256).hexdigest()
+        )
+        if not hmac.compare_digest(computed_sig, expected_sig):
+            raise PermissionDenied("Invalid signature")
+
+    def post(self, request, *args, **kwargs):
+        self.validate_signature(request)
+        external_owner_id = request.data.get("external_owner_id")
+        repo_service_id = request.data.get("repo_service_id")
+        if not external_owner_id or not repo_service_id:
+            return Response("Missing required parameters", status=400)
+        try:
+            owner = Owner.objects.get(service_id=external_owner_id)
+        except Owner.DoesNotExist:
+            raise NotFound("Owner not found")
+
+        is_authorized = True
+
+        app_install = GithubAppInstallation.objects.filter(
+            owner_id=owner.ownerid, app_id=AI_FEATURES_GH_APP_ID
+        ).first()
+
+        if not app_install:
+            is_authorized = False
+
+        else:
+            repo_ids = app_install.repository_service_ids
+            if repo_ids and repo_service_id not in repo_ids:
+                is_authorized = False
+
+        return Response({"is_valid": is_authorized})

api/internal/owner/serializers.py

@@ -7,14 +7,12 @@
 from rest_framework import serializers
 from rest_framework.exceptions import PermissionDenied
 from shared.plan.constants import (
-    PAID_PLANS,
-    SENTRY_PAID_USER_PLAN_REPRESENTATIONS,
     TEAM_PLAN_MAX_USERS,
-    TEAM_PLAN_REPRESENTATIONS,
+    TierName,
 )
 from shared.plan.service import PlanService
 
-from codecov_auth.models import Owner
+from codecov_auth.models import Owner, Plan
 from services.billing import BillingService
 from services.sentry import send_user_webhook as send_sentry_webhook
 
@@ -109,8 +107,14 @@ class StripeCardSerializer(serializers.Serializer):
     last4 = serializers.CharField()
 
 
+class StripeUSBankAccountSerializer(serializers.Serializer):
+    bank_name = serializers.CharField()
+    last4 = serializers.CharField()
+
+
 class StripePaymentMethodSerializer(serializers.Serializer):
     card = StripeCardSerializer(read_only=True)
+    us_bank_account = StripeUSBankAccountSerializer(read_only=True)
     billing_details = serializers.JSONField(read_only=True)
 
 
@@ -131,11 +135,6 @@ def validate_value(self, value: str) -> str:
             plan["value"] for plan in plan_service.available_plans(current_owner)
         ]
         if value not in plan_values:
-            if value in SENTRY_PAID_USER_PLAN_REPRESENTATIONS:
-                log.warning(
-                    "Non-Sentry user attempted to transition to Sentry plan",
-                    extra=dict(owner_id=current_owner.pk, plan=value),
-                )
             raise serializers.ValidationError(
                 f"Invalid value for plan: {value}; must be one of {plan_values}"
             )
@@ -148,8 +147,17 @@ def validate(self, plan: Dict[str, Any]) -> Dict[str, Any]:
                 detail="You cannot update your plan manually, for help or changes to plan, connect with [email protected]"
             )
 
+        active_plans = Plan.objects.select_related("tier").filter(
+            paid_plan=True, is_active=True
+        )
+
+        active_plan_names = set(active_plans.values_list("name", flat=True))
+        team_tier_plans = active_plans.filter(
+            tier__tier_name=TierName.TEAM.value
+        ).values_list("name", flat=True)
+
         # Validate quantity here because we need access to whole plan object
-        if plan["value"] in PAID_PLANS:
+        if plan["value"] in active_plan_names:
             if "quantity" not in plan:
                 raise serializers.ValidationError(
                     "Field 'quantity' required for updating to paid plans"
@@ -178,7 +186,7 @@ def validate(self, plan: Dict[str, Any]) -> Dict[str, Any]:
                     "Quantity or plan for paid plan must be different from the existing one"
                 )
             if (
-                plan["value"] in TEAM_PLAN_REPRESENTATIONS
+                plan["value"] in team_tier_plans
                 and plan["quantity"] > TEAM_PLAN_MAX_USERS
             ):
                 raise serializers.ValidationError(
@@ -213,7 +221,7 @@ def get_plan(self, phase: Dict[str, Any]) -> str:
         plan_name = list(stripe_plan_dict.keys())[
             list(stripe_plan_dict.values()).index(plan_id)
         ]
-        marketing_plan_name = PAID_PLANS[plan_name].billing_rate
+        marketing_plan_name = Plan.objects.get(name=plan_name).marketing_name
         return marketing_plan_name
 
     def get_quantity(self, phase: Dict[str, Any]) -> int:
@@ -336,7 +344,11 @@ def update(self, instance: Owner, validated_data: Dict[str, Any]) -> object:
                 instance, desired_plan
             )
 
-            if desired_plan["value"] in SENTRY_PAID_USER_PLAN_REPRESENTATIONS:
+            sentry_plans = Plan.objects.filter(
+                tier__tier_name=TierName.SENTRY.value, is_active=True
+            ).values_list("name", flat=True)
+
+            if desired_plan["value"] in sentry_plans:
                 current_owner = self.context["view"].request.current_owner
                 send_sentry_webhook(current_owner, instance)
 

api/internal/owner/views.py

@@ -80,12 +80,33 @@ def update_payment(self, request, *args, **kwargs):
     @action(detail=False, methods=["patch"])
     @stripe_safe
     def update_email(self, request, *args, **kwargs):
+        """
+        Update the email address associated with the owner's billing account.
+
+        Args:
+            request: The HTTP request object containing:
+                - new_email: The new email address to update to
+                - apply_to_default_payment_method: Boolean flag to update email on the default payment method (default False)
+
+        Returns:
+            Response with serialized owner data
+
+        Raises:
+            ValidationError: If no new_email is provided in the request
+        """
         new_email = request.data.get("new_email")
         if not new_email:
             raise ValidationError(detail="No new_email sent")
         owner = self.get_object()
         billing = BillingService(requesting_user=request.current_owner)
-        billing.update_email_address(owner, new_email)
+        apply_to_default_payment_method = request.data.get(
+            "apply_to_default_payment_method", False
+        )
+        billing.update_email_address(
+            owner,
+            new_email,
+            apply_to_default_payment_method=apply_to_default_payment_method,
+        )
         return Response(self.get_serializer(owner).data)
 
     @action(detail=False, methods=["patch"])

api/internal/tests/test_pagination.py

@@ -1,14 +1,18 @@
 from rest_framework.reverse import reverse
 from rest_framework.test import APITestCase
+from shared.django_apps.codecov_auth.tests.factories import PlanFactory, TierFactory
 from shared.django_apps.core.tests.factories import OwnerFactory
+from shared.plan.constants import TierName
 
 from utils.test_utils import Client
 
 
 class PageNumberPaginationTests(APITestCase):
     def setUp(self):
         self.client = Client()
-        self.owner = OwnerFactory(plan="users-free", plan_user_count=5)
+        tier = TierFactory(tier_name=TierName.BASIC.value)
+        plan = PlanFactory(tier=tier, is_active=True)
+        self.owner = OwnerFactory(plan=plan.name, plan_user_count=5)
         self.users = [
             OwnerFactory(organizations=[self.owner.ownerid]),
             OwnerFactory(organizations=[self.owner.ownerid]),