ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-23 04:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Django tests

Browse Source
This commit is contained in:
Chandrapal Badshah
2025-04-29 19:36:21 +05:30
parent e70357fba1
commit c459cde1ff
11 changed files with 298 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/backend/api/exceptions.py
+1 -1
View File
@@ -3,7 +3,7 @@ from rest_framework import status
from rest_framework.exceptions import APIException
from rest_framework_json_api.exceptions import exception_handler
from rest_framework_json_api.serializers import ValidationError
from rest_framework_simplejwt.exceptions import TokenError, InvalidToken
from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
class ModelValidationError(ValidationError):
api/src/backend/api/migrations/0017_lighthouseconfig.py
+4 -2
View File
@@ -1,10 +1,12 @@
# Generated by Django 5.1.7 on 2025-04-10 14:54
import api.rls
import uuid
import django.core.validators
import django.db.models.deletion
import uuid
from django.db import migrations, models
import api.rls
class Migration(migrations.Migration):
dependencies = [
api/src/backend/api/models.py
+89 -7
View File
@@ -1,7 +1,8 @@
import json
import logging
import re
from uuid import UUID, uuid4
import logging
from cryptography.fernet import Fernet, InvalidToken
from django.conf import settings
from django.contrib.auth.models import AbstractBaseUser
@@ -1329,15 +1330,39 @@ class LighthouseConfig(RowLevelSecurityProtectedModel):
Stores configuration and API keys for LLM services.
"""
MODEL_CHOICES = [
"gpt-4o-2024-11-20",
"gpt-4o-2024-08-06",
"gpt-4o-2024-05-13",
"gpt-4o",
"gpt-4o-mini-2024-07-18",
"gpt-4o-mini",
]
id = models.UUIDField(primary_key=True, default=uuid4, editable=False)
inserted_at = models.DateTimeField(auto_now_add=True, editable=False)
updated_at = models.DateTimeField(auto_now=True, editable=False)
name = models.CharField(max_length=100, validators=[MinLengthValidator(3)])
api_key = models.BinaryField(blank=False, null=False)
model = models.CharField(max_length=50, default="gpt-4o-2024-08-06")
temperature = models.FloatField(default=0)
max_tokens = models.IntegerField(default=4000)
name = models.CharField(
max_length=100,
validators=[MinLengthValidator(3)],
blank=False,
null=False,
help_text="Name of the configuration",
)
api_key = models.BinaryField(
blank=False, null=False, help_text="Encrypted API key for the LLM service"
)
model = models.CharField(
max_length=50,
blank=False,
null=False,
help_text="Must be one of the supported model names",
)
temperature = models.FloatField(default=0, help_text="Must be between 0 and 1")
max_tokens = models.IntegerField(
default=4000, help_text="Must be between 500 and 5000"
)
business_context = models.TextField(
blank=True,
null=True,
@@ -1348,6 +1373,33 @@ class LighthouseConfig(RowLevelSecurityProtectedModel):
def __str__(self):
return self.name
def clean(self):
super().clean()
# Validate temperature
if not 0 <= self.temperature <= 1:
raise ModelValidationError(
detail="Temperature must be between 0 and 1",
code="invalid_temperature",
pointer="/data/attributes/temperature",
)
# Validate max_tokens
if not 500 <= self.max_tokens <= 5000:
raise ModelValidationError(
detail="Max tokens must be between 500 and 5000",
code="invalid_max_tokens",
pointer="/data/attributes/max_tokens",
)
# Validate model name
if self.model not in self.MODEL_CHOICES:
raise ModelValidationError(
detail=f"Model must be one of: {', '.join(self.MODEL_CHOICES)}",
code="invalid_model",
pointer="/data/attributes/model",
)
@property
def api_key_decoded(self):
"""Return the decrypted API key."""
@@ -1369,12 +1421,42 @@ class LighthouseConfig(RowLevelSecurityProtectedModel):
@api_key_decoded.setter
def api_key_decoded(self, value):
"""Store the encrypted API key."""
if not value:
raise ModelValidationError(
detail="API key is required",
code="invalid_api_key",
pointer="/data/attributes/api_key",
)
# Validate OpenAI API key format
openai_key_pattern = r"^sk-[\w-]+T3BlbkFJ[\w-]+$"
if not re.match(openai_key_pattern, value):
raise ValueError("Invalid OpenAI API key format. Must start with 'sk-'.")
raise ModelValidationError(
detail="Invalid OpenAI API key format. Must start with 'sk-'.",
code="invalid_api_key",
pointer="/data/attributes/api_key",
)
self.api_key = fernet.encrypt(value.encode())
def save(self, *args, **kwargs):
# Validate required fields
if not self.name:
raise ModelValidationError(
detail="Name is required",
code="missing_name",
pointer="/data/attributes/name",
)
if not self.model:
raise ModelValidationError(
detail="Model is required",
code="missing_model",
pointer="/data/attributes/model",
)
self.full_clean()
super().save(*args, **kwargs)
class Meta(RowLevelSecurityProtectedModel.Meta):
db_table = "lighthouse_config"
api/src/backend/api/partitions.py
+1 -1
View File
@@ -4,11 +4,11 @@ from typing import Generator, Optional
from dateutil.relativedelta import relativedelta
from django.conf import settings
from psqlextra.partitioning import (
PostgresPartitioningError,
PostgresPartitioningManager,
PostgresRangePartition,
PostgresRangePartitioningStrategy,
PostgresTimePartitionSize,
PostgresPartitioningError,
)
from psqlextra.partitioning.config import PostgresPartitioningConfig
from uuid6 import UUID
api/src/backend/api/tests/integration/test_tenants.py
+1 -2
View File
@@ -1,10 +1,9 @@
from unittest.mock import patch
import pytest
from conftest import TEST_PASSWORD, TEST_USER, get_api_tokens, get_authorization_header
from django.urls import reverse
from conftest import TEST_USER, TEST_PASSWORD, get_api_tokens, get_authorization_header
@patch("api.v1.views.schedule_provider_scan")
@pytest.mark.django_db
api/src/backend/api/tests/test_compliance.py
+5 -5
View File
@@ -1,12 +1,12 @@
from unittest.mock import patch, MagicMock
from unittest.mock import MagicMock, patch
from api.compliance import (
generate_compliance_overview_template,
generate_scan_compliance,
get_prowler_provider_checks,
get_prowler_provider_compliance,
load_prowler_compliance,
load_prowler_checks,
generate_scan_compliance,
generate_compliance_overview_template,
load_prowler_compliance,
)
from api.models import Provider
@@ -69,7 +69,7 @@ class TestCompliance:
load_prowler_compliance()
from api.compliance import PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE, PROWLER_CHECKS
from api.compliance import PROWLER_CHECKS, PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE
assert PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE == {
"template_key": "template_value"
api/src/backend/api/tests/test_database.py
+3 -2
View File
@@ -1,12 +1,13 @@
from unittest.mock import patch
import pytest
from config.django.base import DATABASE_ROUTERS as PROD_DATABASE_ROUTERS
from django.conf import settings
from django.db.migrations.recorder import MigrationRecorder
from django.db.utils import ConnectionRouter
from api.db_router import MainRouter
from api.rls import Tenant
from config.django.base import DATABASE_ROUTERS as PROD_DATABASE_ROUTERS
from unittest.mock import patch
@patch("api.db_router.MainRouter.admin_db", new="admin")
api/src/backend/api/tests/test_uuid_utils.py
+3 -3
View File
@@ -7,12 +7,12 @@ from rest_framework_json_api.serializers import ValidationError
from uuid6 import UUID
from api.uuid_utils import (
transform_into_uuid7,
datetime_to_uuid7,
datetime_from_uuid7,
uuid7_start,
datetime_to_uuid7,
transform_into_uuid7,
uuid7_end,
uuid7_range,
uuid7_start,
)
api/src/backend/api/tests/test_views.py
+176 -123
View File
@@ -5317,7 +5317,7 @@ class TestLighthouseConfigViewSet:
"data": {
"type": "lighthouse-config",
"attributes": {
"name": "Test Config",
"name": "OpenAI",
"api_key": "sk-test1234567890T3BlbkFJtest1234567890",
"model": "gpt-4o",
"temperature": 0.7,
@@ -5344,34 +5344,152 @@ class TestLighthouseConfigViewSet:
}
def test_lighthouse_config_list(self, authenticated_client):
response = authenticated_client.get("/api/v1/lighthouse-config/")
response = authenticated_client.get(reverse("lighthouseconfig-list"))
assert response.status_code == status.HTTP_200_OK
assert response.json()["data"] == []
def test_lighthouse_config_create(self, authenticated_client, valid_config_payload):
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
reverse("lighthouseconfig-list"),
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
data = response.json()["data"]
assert data["attributes"]["name"] == "Test Config"
assert data["attributes"]["model"] == "gpt-4o"
assert data["attributes"]["temperature"] == 0.7
assert data["attributes"]["max_tokens"] == 4000
assert data["attributes"]["business_context"] == "Test business context"
assert data["attributes"]["is_active"] is True
assert data["attributes"]["api_key"] == "*" * len(
valid_config_payload["data"]["attributes"]["api_key"]
assert (
data["attributes"]["name"]
== valid_config_payload["data"]["attributes"]["name"]
)
assert (
data["attributes"]["model"]
== valid_config_payload["data"]["attributes"]["model"]
)
assert (
data["attributes"]["temperature"]
== valid_config_payload["data"]["attributes"]["temperature"]
)
assert (
data["attributes"]["max_tokens"]
== valid_config_payload["data"]["attributes"]["max_tokens"]
)
assert (
data["attributes"]["business_context"]
== valid_config_payload["data"]["attributes"]["business_context"]
)
assert (
data["attributes"]["is_active"]
== valid_config_payload["data"]["attributes"]["is_active"]
)
# Check that API key is masked with asterisks only
masked_api_key = data["attributes"]["api_key"]
assert all(
c == "*" for c in masked_api_key
), "API key should contain only asterisks"
def test_lighthouse_config_create_invalid_name_too_short(
self, authenticated_client, valid_config_payload
):
"""Test that name validation fails when too short"""
payload = valid_config_payload.copy()
payload["data"]["attributes"]["name"] = "T" # Too short
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("name" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_invalid_api_key_format(
self, authenticated_client, valid_config_payload
):
"""Test that API key validation fails with invalid format"""
payload = valid_config_payload.copy()
payload["data"]["attributes"]["api_key"] = "invalid-key"
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("api_key" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_invalid_model(
self, authenticated_client, valid_config_payload
):
"""Test that model validation fails with invalid model name"""
payload = valid_config_payload.copy()
payload["data"]["attributes"]["model"] = "invalid-model"
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("model" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_invalid_temperature_range(
self, authenticated_client, valid_config_payload
):
"""Test that temperature validation fails when out of range"""
payload = valid_config_payload.copy()
payload["data"]["attributes"]["temperature"] = 2.0 # Out of range
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("temperature" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_invalid_max_tokens(
self, authenticated_client, valid_config_payload
):
"""Test that max_tokens validation fails with invalid value"""
payload = valid_config_payload.copy()
payload["data"]["attributes"]["max_tokens"] = -1 # Invalid value
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("max_tokens" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_missing_required_fields(
self, authenticated_client
):
"""Test that validation fails when required fields are missing"""
payload = {"data": {"type": "lighthouse-config", "attributes": {}}}
response = authenticated_client.post(
reverse("lighthouseconfig-list"),
data=payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
# Check for required fields
required_fields = ["name", "api_key", "model"]
for field in required_fields:
assert any(field in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_create_duplicate(
self, authenticated_client, valid_config_payload
):
# Create first config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
reverse("lighthouseconfig-list"),
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
@@ -5379,7 +5497,7 @@ class TestLighthouseConfigViewSet:
# Try to create second config for same tenant
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
reverse("lighthouseconfig-list"),
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
@@ -5389,55 +5507,29 @@ class TestLighthouseConfigViewSet:
in response.json()["errors"][0]["detail"]
)
def test_lighthouse_config_create_invalid(
self, authenticated_client, invalid_config_payload
):
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=invalid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_400_BAD_REQUEST
errors = response.json()["errors"]
assert any("name" in error["source"]["pointer"] for error in errors)
assert any("api_key" in error["source"]["pointer"] for error in errors)
def test_lighthouse_config_retrieve(
self, authenticated_client, valid_config_payload
self, authenticated_client, lighthouse_config_fixture
):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
"""Test retrieving a lighthouse config"""
response = authenticated_client.get(
reverse(
"lighthouseconfig-detail", kwargs={"pk": lighthouse_config_fixture.id}
)
)
assert response.status_code == status.HTTP_201_CREATED
config_id = response.json()["data"]["id"]
# Retrieve config
response = authenticated_client.get(f"/api/v1/lighthouse-config/{config_id}/")
assert response.status_code == status.HTTP_200_OK
data = response.json()["data"]
assert data["attributes"]["name"] == "Test Config"
assert data["attributes"]["name"] == lighthouse_config_fixture.name
assert data["attributes"]["api_key"] == "*" * len(
valid_config_payload["data"]["attributes"]["api_key"]
lighthouse_config_fixture.api_key
)
def test_lighthouse_config_update(self, authenticated_client, valid_config_payload):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
config_id = response.json()["data"]["id"]
# Update config
def test_lighthouse_config_update(
self, authenticated_client, lighthouse_config_fixture
):
update_payload = {
"data": {
"type": "lighthouse-config",
"id": config_id,
"id": str(lighthouse_config_fixture.id),
"attributes": {
"name": "Updated Config",
"model": "gpt-4o-mini",
@@ -5446,7 +5538,9 @@ class TestLighthouseConfigViewSet:
}
}
response = authenticated_client.patch(
f"/api/v1/lighthouse-config/{config_id}/",
reverse(
"lighthouseconfig-detail", kwargs={"pk": lighthouse_config_fixture.id}
),
data=update_payload,
content_type=API_JSON_CONTENT_TYPE,
)
@@ -5456,128 +5550,87 @@ class TestLighthouseConfigViewSet:
assert data["attributes"]["model"] == "gpt-4o-mini"
assert data["attributes"]["temperature"] == 0.5
def test_lighthouse_config_delete(self, authenticated_client, valid_config_payload):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
config_id = response.json()["data"]["id"]
# Delete config
def test_lighthouse_config_delete(
self, authenticated_client, lighthouse_config_fixture
):
config_id = lighthouse_config_fixture.id
response = authenticated_client.delete(
f"/api/v1/lighthouse-config/{config_id}/"
reverse("lighthouseconfig-detail", kwargs={"pk": config_id})
)
assert response.status_code == status.HTTP_204_NO_CONTENT
# Verify deletion
response = authenticated_client.get(f"/api/v1/lighthouse-config/{config_id}/")
response = authenticated_client.get(
reverse("lighthouseconfig-detail", kwargs={"pk": config_id})
)
assert response.status_code == status.HTTP_404_NOT_FOUND
@patch("api.v1.views.openai.OpenAI")
def test_lighthouse_config_check_connection(
self, mock_openai, authenticated_client, valid_config_payload
self, mock_openai, authenticated_client, lighthouse_config_fixture
):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
config_id = response.json()["data"]["id"]
config_id = lighthouse_config_fixture.id
# Mock successful API call
mock_client = Mock()
mock_client.models.list.return_value = Mock(
data=[Mock(id="gpt-4o"), Mock(id="gpt-4o-mini")]
)
mock_openai.return_value = mock_client
# Check connection
response = authenticated_client.get(
f"/api/v1/lighthouse-config/{config_id}/check_connection/"
reverse("lighthouseconfig-check-connection", kwargs={"pk": config_id})
)
assert response.status_code == status.HTTP_200_OK
assert response.json()["detail"] == "Connection successful!"
assert "gpt-4o" in response.json()["available_models"]
assert "gpt-4o-mini" in response.json()["available_models"]
assert response.json()["data"]["detail"] == "Connection successful!"
assert "gpt-4o" in response.json()["data"]["available_models"]
assert "gpt-4o-mini" in response.json()["data"]["available_models"]
def test_lighthouse_config_show_key(
self, authenticated_client, valid_config_payload
self, authenticated_client, lighthouse_config_fixture, valid_config_payload
):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
config_id = response.json()["data"]["id"]
# Show key
config_id = lighthouse_config_fixture.id
expected_api_key = valid_config_payload["data"]["attributes"]["api_key"]
response = authenticated_client.get(
f"/api/v1/lighthouse-config/{config_id}/show_key/"
reverse("lighthouseconfig-show-key", kwargs={"pk": config_id})
)
assert response.status_code == status.HTTP_200_OK
assert (
response.json()["api_key"]
== valid_config_payload["data"]["attributes"]["api_key"]
)
assert response.json()["data"]["api_key"] == expected_api_key
def test_lighthouse_config_filters(
self, authenticated_client, valid_config_payload
self, authenticated_client, lighthouse_config_fixture
):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
# Test name filter
response = authenticated_client.get(
"/api/v1/lighthouse-config/?filter[name]=Test Config"
reverse("lighthouseconfig-list")
+ "?filter[name]="
+ lighthouse_config_fixture.name
)
assert response.status_code == status.HTTP_200_OK
assert len(response.json()["data"]) == 1
# Test model filter
response = authenticated_client.get(
"/api/v1/lighthouse-config/?filter[model]=gpt-4o"
reverse("lighthouseconfig-list") + "?filter[model]=gpt-4o"
)
assert response.status_code == status.HTTP_200_OK
assert len(response.json()["data"]) == 1
# Test is_active filter
response = authenticated_client.get(
"/api/v1/lighthouse-config/?filter[is_active]=true"
reverse("lighthouseconfig-list") + "?filter[is_active]=true"
)
assert response.status_code == status.HTTP_200_OK
assert len(response.json()["data"]) == 1
def test_lighthouse_config_sorting(
self, authenticated_client, valid_config_payload
self, authenticated_client, lighthouse_config_fixture
):
# Create config
response = authenticated_client.post(
"/api/v1/lighthouse-config/",
data=valid_config_payload,
content_type=API_JSON_CONTENT_TYPE,
)
assert response.status_code == status.HTTP_201_CREATED
# Test sorting by name
response = authenticated_client.get("/api/v1/lighthouse-config/?sort=name")
response = authenticated_client.get(
reverse("lighthouseconfig-list") + "?sort=name"
)
assert response.status_code == status.HTTP_200_OK
assert len(response.json()["data"]) == 1
# Test sorting by inserted_at
response = authenticated_client.get(
"/api/v1/lighthouse-config/?sort=-inserted_at"
reverse("lighthouseconfig-list") + "?sort=-inserted_at"
)
assert response.status_code == status.HTTP_200_OK
assert len(response.json()["data"]) == 1
api/src/backend/api/v1/urls.py
+1 -1
View File
@@ -3,7 +3,6 @@ from drf_spectacular.views import SpectacularRedocView
from rest_framework_nested import routers
from api.v1.views import (
LighthouseConfigViewSet,
ComplianceOverviewViewSet,
CustomTokenObtainView,
CustomTokenRefreshView,
@@ -14,6 +13,7 @@ from api.v1.views import (
IntegrationViewSet,
InvitationAcceptViewSet,
InvitationViewSet,
LighthouseConfigViewSet,
MembershipViewSet,
OverviewViewSet,
ProviderGroupProvidersRelationshipView,
api/src/backend/conftest.py
+14
View File
@@ -19,6 +19,7 @@ from api.models import (
Integration,
IntegrationProviderRelationship,
Invitation,
LighthouseConfig,
Membership,
Provider,
ProviderGroup,
@@ -928,6 +929,19 @@ def backfill_scan_metadata_fixture(scans_fixture, findings_fixture):
scan_id = scan_instance.id
backfill_resource_scan_summaries(tenant_id=tenant_id, scan_id=scan_id)
@pytest.fixture
def lighthouse_config_fixture(authenticated_client, tenants_fixture):
return LighthouseConfig.objects.create(
tenant_id=tenants_fixture[0].id,
name="OpenAI",
api_key_decoded="sk-test1234567890T3BlbkFJtest1234567890",
model="gpt-4o",
temperature=0,
max_tokens=4000,
business_context="Test business context",
is_active=True,
)
def get_authorization_header(access_token: str) -> dict:
return {"Authorization": f"Bearer {access_token}"}