Fix Live Evaluation API tests #1953

Signed-off-by: Michael Ehab Mikhail <michael.ehab@hotmail.com>

Author: michaelehab

vulnerabilities/api_v2.py

@@ -44,6 +44,7 @@
 from vulnerabilities.models import Weakness
 from vulnerabilities.tasks import enqueue_ad_hoc_pipeline
 from vulnerabilities.throttling import PermissionBasedUserRateThrottle
+from vulnerablecode.settings import VULNERABLECODE_ENABLE_LIVE_EVALUATION_API
 
 
 class CharInFilter(filters.BaseInFilter, filters.CharFilter):
@@ -1316,6 +1317,11 @@ class LiveEvaluationViewSet(viewsets.GenericViewSet):
     )
     @action(detail=False, methods=["post"])
     def evaluate(self, request):
+        if not VULNERABLECODE_ENABLE_LIVE_EVALUATION_API:
+            return Response(
+                {"error": "Live evaluation API is disabled."},
+                status=status.HTTP_403_FORBIDDEN,
+            )
         serializer = self.get_serializer(data=request.data)
         if not serializer.is_valid():
             return Response(
@@ -1385,6 +1391,12 @@ def evaluate(self, request):
     )
     @action(detail=False, methods=["get"], url_path=r"status/(?P<live_run_id>[0-9a-f\-]{36})")
     def status(self, request, live_run_id=None):
+        if not VULNERABLECODE_ENABLE_LIVE_EVALUATION_API:
+            return Response(
+                {"error": "Live evaluation API is disabled."},
+                status=status.HTTP_403_FORBIDDEN,
+            )
+
         from vulnerabilities.models import LivePipelineRun
 
         try:

vulnerabilities/tests/test_api_v2.py

@@ -912,6 +912,7 @@ def setUp(self):
         self.client = APIClient(enforce_csrf_checks=True)
         self.url = "/api/v2/live-evaluation/evaluate"
 
+    @patch("vulnerabilities.api_v2.VULNERABLECODE_ENABLE_LIVE_EVALUATION_API", True)
     @patch("vulnerabilities.api_v2.LIVE_IMPORTERS_REGISTRY")
     @patch("vulnerabilities.api_v2.enqueue_ad_hoc_pipeline")
     @patch("django.urls.reverse")
@@ -920,7 +921,6 @@ class MockImporter:
             pipeline_id = "pypa_live_importer_v2"
             supported_types = ["pypi"]
 
-        os.environ["VULNERABLECODE_ENABLE_LIVE_EVALUATION_API"] = "true"
         mock_registry.values.return_value = [MockImporter]
         valid_uuid = "00000000-0000-0000-0000-000000000001"
         mock_enqueue.return_value = (valid_uuid, ["mock-run-id"])
@@ -936,30 +936,36 @@ class MockImporter:
         assert "status_url" in response.data
         assert response.data["status_url"].endswith(f"/api/v2/live-evaluation/status/{valid_uuid}")
 
+    @patch("vulnerabilities.api_v2.VULNERABLECODE_ENABLE_LIVE_EVALUATION_API", True)
     @patch("vulnerabilities.api_v2.LIVE_IMPORTERS_REGISTRY")
     def test_evaluate_no_importer_found(self, mock_registry):
         class MockImporter:
             pipeline_id = "dummy"
             supported_types = ["npm"]
 
-        os.environ["VULNERABLECODE_ENABLE_LIVE_EVALUATION_API"] = "true"
         mock_registry.values.return_value = [MockImporter]
         data = {"purl": "pkg:pypi/django@3.2"}
         response = self.client.post(self.url, data, format="json")
         assert response.status_code == 400
         assert "No live importers found" in response.data["error"]
 
+    @patch("vulnerabilities.api_v2.VULNERABLECODE_ENABLE_LIVE_EVALUATION_API", True)
     def test_evaluate_invalid_purl(self):
-        os.environ["VULNERABLECODE_ENABLE_LIVE_EVALUATION_API"] = "true"
         data = {"purl": "not_a_valid_purl"}
         response = self.client.post(self.url, data, format="json")
         assert response.status_code == 400
         assert "Invalid PackageURL" in response.data["error"]
 
+    @patch("vulnerabilities.api_v2.VULNERABLECODE_ENABLE_LIVE_EVALUATION_API", True)
     @patch("vulnerabilities.models.LivePipelineRun.objects.get")
     def test_status_not_found(self, mock_live_get):
-        os.environ["VULNERABLECODE_ENABLE_LIVE_EVALUATION_API"] = "true"
         mock_live_get.side_effect = LivePipelineRun.DoesNotExist()
         url = "/api/v2/live-evaluation/status/00000000-0000-0000-0000-000000000000"
         response = self.client.get(url)
         assert response.status_code == 404
+
+    @patch("vulnerabilities.api_v2.VULNERABLECODE_ENABLE_LIVE_EVALUATION_API", False)
+    def test_evaluate_disabled_returns_403(self):
+        data = {"purl": "pkg:pypi/django@3.2"}
+        response = self.client.post(self.url, data, format="json")
+        assert response.status_code == 403

vulnerablecode/urls.py

@@ -46,7 +46,6 @@
 from vulnerabilities.views import VulnerabilitySearch
 from vulnerablecode.settings import DEBUG
 from vulnerablecode.settings import DEBUG_TOOLBAR
-from vulnerablecode.settings import VULNERABLECODE_ENABLE_LIVE_EVALUATION_API
 
 
 # See the comment at https://stackoverflow.com/a/46163870.
@@ -73,8 +72,7 @@ def __init__(self, *args, **kwargs):
 api_v2_router.register("pipelines", PipelineScheduleV2ViewSet, basename="pipelines")
 api_v2_router.register("advisory-codefixes", CodeFixV2ViewSet, basename="advisory-codefix")
 
-if VULNERABLECODE_ENABLE_LIVE_EVALUATION_API:
-    api_v2_router.register("live-evaluation", LiveEvaluationViewSet, basename="live-evaluation")
+api_v2_router.register("live-evaluation", LiveEvaluationViewSet, basename="live-evaluation")
 
 
 urlpatterns = [