Lint

Author: rohitvinnakota-codecov

api/gen_ai/tests/test_gen_ai.py

@@ -1,24 +1,25 @@
-import hmac, json
+import hmac
+import json
 from hashlib import sha256
 from unittest.mock import patch
 
 from django.urls import reverse
-from django.utils.crypto import constant_time_compare
 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 Owner, GithubAppInstallation
-from utils.config import get_config
-from shared.django_apps.core.tests.factories import OwnerFactory, RepositoryFactory
+from codecov_auth.models import GithubAppInstallation
 
 PAYLOAD_SECRET = b"testixik8qdauiab1yiffydimvi72ekq"
 VIEW_URL = reverse("auth")
 
+
 def sign_payload(payload, secret=PAYLOAD_SECRET):
     data = json.dumps(payload, separators=(",", ":")).encode("utf-8")
     signature = "sha256=" + hmac.new(secret, data, digestmod=sha256).hexdigest()
     return signature, data
 
+
 class GenAIAuthViewTests(APITestCase):
     @patch("utils.config.get_config", return_value=PAYLOAD_SECRET)
     def test_missing_parameters(self, mock_config):
@@ -37,7 +38,9 @@ def test_missing_parameters(self, mock_config):
     def test_invalid_signature(self, mock_config):
         payload = {"external_owner_id": "owner1", "repo_service_id": "101"}
         # Create a wrong signature by altering the payload before signing
-        wrong_sig = "sha256=" + hmac.new(PAYLOAD_SECRET, b"{}", digestmod=sha256).hexdigest()
+        wrong_sig = (
+            "sha256=" + hmac.new(PAYLOAD_SECRET, b"{}", digestmod=sha256).hexdigest()
+        )
         response = self.client.post(
             VIEW_URL,
             data=payload,
@@ -48,21 +51,20 @@ def test_invalid_signature(self, mock_config):
 
     @patch("utils.config.get_config", return_value=PAYLOAD_SECRET)
     def test_owner_not_found(self, mock_config):
-        payload = {'external_owner_id': 'nonexistent_owner', 'repo_service_id': '101'}
+        payload = {"external_owner_id": "nonexistent_owner", "repo_service_id": "101"}
         sig, serialized_data = sign_payload(payload)
         response = self.client.post(
             VIEW_URL,
             HTTP_X_GEN_AI_AUTH_SIGNATURE=sig,
             data=serialized_data,
             content_type="application/json",
-            
         )
         self.assertEqual(response.status_code, 404)
 
     @patch("utils.config.get_config", return_value=PAYLOAD_SECRET)
     def test_no_installation(self, mock_config):
         _ = OwnerFactory(service="github", service_id="owner1", username="test1")
-        payload = {'external_owner_id': 'owner1', 'repo_service_id': '101'}
+        payload = {"external_owner_id": "owner1", "repo_service_id": "101"}
         sig, data = sign_payload(payload)
         response = self.client.post(
             VIEW_URL,
@@ -81,10 +83,10 @@ def test_authorized(self, mock_config):
             installation_id=12345,
             owner=owner,
             name="ai-features",
-            repository_service_ids=['101', '202']
+            repository_service_ids=["101", "202"],
         )
         app_install.save()
-        payload = {"external_owner_id": "owner2", "repo_service_id": '101'}
+        payload = {"external_owner_id": "owner2", "repo_service_id": "101"}
         sig, data = sign_payload(payload)
         response = self.client.post(
             VIEW_URL,
@@ -103,10 +105,10 @@ def test_unauthorized(self, mock_config):
             installation_id=2,
             owner=owner,
             name="ai-features",
-            repository_service_ids=["303", "404"]
+            repository_service_ids=["303", "404"],
         )
         app_install.save()
-        payload = {'external_owner_id': 'owner3', 'repo_service_id': '101'}
+        payload = {"external_owner_id": "owner3", "repo_service_id": "101"}
         sig, data = sign_payload(payload)
         response = self.client.post(
             VIEW_URL,

api/gen_ai/views.py

@@ -1,19 +1,20 @@
 import hmac
 import logging
 from hashlib import sha256
+
 from django.utils.crypto import constant_time_compare
-from graphql_api.types.owner.owner import AI_FEATURES_GH_APP_ID
 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 codecov_auth.models import Owner
-from core.models import Repository
+
+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
-from codecov_auth.models import GithubAppInstallation
 
 log = logging.getLogger(__name__)
 
+
 class GenAIAuthView(APIView):
     permission_classes = [AllowAny]
 
@@ -24,7 +25,9 @@ def validate_signature(self, request):
         if isinstance(key, str):
             key = key.encode("utf-8")
         expected_sig = request.META.get("HTTP_X_GEN_AI_AUTH_SIGNATURE")
-        computed_sig = "sha256=" + hmac.new(key, request.body, digestmod=sha256).hexdigest()
+        computed_sig = (
+            "sha256=" + hmac.new(key, request.body, digestmod=sha256).hexdigest()
+        )
         if not (expected_sig and constant_time_compare(computed_sig, expected_sig)):
             raise PermissionDenied("Invalid signature")
 
@@ -38,23 +41,24 @@ def post(self, request, *args, **kwargs):
             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()
+        app_install = GithubAppInstallation.objects.filter(
+            owner_id=owner.ownerid, app_id=AI_FEATURES_GH_APP_ID
+        ).first()
 
         if not app_install:
             print("FAILED")
             is_authorized = False
-        
+
         else:
             repo_ids = app_install.repository_service_ids
             if repo_ids and repo_service_id not in repo_ids:
                 print("HERE")
                 is_authorized = False
 
-
         return Response({"is_valid": is_authorized})
 
 
-# api/gen_ai/tests/test_gen_ai.py 
+# api/gen_ai/tests/test_gen_ai.py