Added graph and table to document
This commit is contained in:
131
backend/core/tests/test_utils.py
Normal file
131
backend/core/tests/test_utils.py
Normal file
@@ -0,0 +1,131 @@
|
||||
from django.test import TestCase
|
||||
from unittest.mock import patch, MagicMock
|
||||
from django.contrib.staticfiles.testing import StaticLiveServerTestCase
|
||||
from io import BytesIO
|
||||
import base64
|
||||
import matplotlib
|
||||
matplotlib.use('Agg')
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from backend.core.models import Organization, Risk, Control, Document
|
||||
from backend.core.utils import *
|
||||
|
||||
class UtilsTests(TestCase):
|
||||
def setUp(self):
|
||||
self.organization = Organization.objects.create(
|
||||
id=1,
|
||||
name="Test Organization",
|
||||
email="test@example.com",
|
||||
employee_headcount="100-500",
|
||||
annual_revenue="$1M-$10M",
|
||||
critical_applications="5-10",
|
||||
compliance_frameworks=["Ab", "Ba"],
|
||||
industry_sector="Technology",
|
||||
it_dependency=8,
|
||||
data_sensitivity="High",
|
||||
network_infrastructure="Cloud-based",
|
||||
remote_workforce_percentage="50%",
|
||||
third_party_vendor_access="10-20",
|
||||
internal_software_development="Moderate",
|
||||
geographic_scope="Global",
|
||||
customer_base="Enterprise",
|
||||
customer_type="B2B",
|
||||
product_portfolio="Diverse",
|
||||
supplier_base="International",
|
||||
it_infrastructure=["Cloud", "On-Premise"],
|
||||
intellectual_property=["Patents", "Trademarks"],
|
||||
sensitive_data=["PII", "Financial Data"],
|
||||
integration_level="Highly Integrated"
|
||||
)
|
||||
|
||||
self.risk = Risk.objects.create(
|
||||
risk_id=1,
|
||||
risk_name="Test Risk",
|
||||
category="Security",
|
||||
primary_impact="Financial"
|
||||
)
|
||||
|
||||
self.controls = [Control.objects.create(id=i, name=f"Control {i}") for i in range(1, 11)]
|
||||
|
||||
def test_extract_organization_details(self):
|
||||
details = extract_organization_details(self.organization)
|
||||
self.assertNotIn('name', details)
|
||||
self.assertNotIn('email', details)
|
||||
self.assertIn("What is your organization's current employee headcount?", details)
|
||||
self.assertEqual(details["What is your organization's current employee headcount?"], "100-500")
|
||||
|
||||
@patch('backend.core.utils.OpenAI')
|
||||
def test_get_top_risk(self, mock_openai):
|
||||
mock_client = MagicMock()
|
||||
mock_openai.return_value = mock_client
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices[0].message.content = "1,2,3"
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
|
||||
risks = get_top_risk(self.organization)
|
||||
self.assertEqual(risks, [1,2,3])
|
||||
|
||||
@patch('backend.core.utils.OpenAI')
|
||||
def test_get_controls_for_risk(self, mock_openai):
|
||||
mock_client = MagicMock()
|
||||
mock_openai.return_value = mock_client
|
||||
mock_response = MagicMock()
|
||||
|
||||
control_lines = [f"{i} : 8 : 5" for i in range(1, 11)]
|
||||
mock_response.choices[0].message.content = "\n".join(control_lines)
|
||||
mock_client.chat.completions.create.return_value = mock_response
|
||||
|
||||
controls = get_controls_for_risk(self.risk, self.organization)
|
||||
self.assertEqual(len(controls), 10)
|
||||
self.assertEqual(controls[0][0], 1)
|
||||
|
||||
@patch('backend.core.utils.HTML')
|
||||
def test_generate_pdf(self, mock_html):
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.write_pdf.return_value = b'PDF_CONTENT'
|
||||
mock_html.return_value = mock_instance
|
||||
|
||||
doc = Document.objects.create(organization=self.organization)
|
||||
response = generate_pdf(doc)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response['Content-Type'], 'application/pdf')
|
||||
|
||||
def test_calculate_aggregate_weight(self):
|
||||
controls = [{'weight': 5}, {'weight': 3}]
|
||||
self.assertEqual(calculate_aggregate_weight(controls), 8)
|
||||
|
||||
def test_calculate_aggregate_likelihood(self):
|
||||
controls = [{'likelihood': 2}, {'likelihood': 4}]
|
||||
self.assertEqual(calculate_aggregate_likelihood(controls), 6)
|
||||
|
||||
def test_map_weight_to_impact_likelihood(self):
|
||||
impact, likelihood = map_weight_to_impact_likelihood(50, 30, 100)
|
||||
self.assertAlmostEqual(impact, 5.0)
|
||||
self.assertAlmostEqual(likelihood, 3.0)
|
||||
|
||||
@patch('pdf2image.convert_from_bytes')
|
||||
@patch('backend.core.utils.HTML')
|
||||
def test_generate_first_page_image(self, mock_html, mock_convert):
|
||||
mock_pdf_instance = MagicMock()
|
||||
mock_pdf_instance.write_pdf.return_value = b'PDF_CONTENT'
|
||||
mock_html.return_value = mock_pdf_instance
|
||||
|
||||
mock_image = MagicMock()
|
||||
mock_convert.return_value = [mock_image]
|
||||
|
||||
doc = Document.objects.create(organization=self.organization)
|
||||
img_io = generate_first_page_image(doc)
|
||||
|
||||
self.assertIsInstance(img_io, BytesIO)
|
||||
mock_convert.assert_called_once_with(b'PDF_CONTENT', first_page=1, last_page=1)
|
||||
|
||||
def test_generate_risk_graph(self):
|
||||
risks_with_controls = [
|
||||
{'risk': {'id': 1}, 'impact': 5.0, 'likelihood': 3.0},
|
||||
{'risk': {'id': 2}, 'impact': 7.0, 'likelihood': 4.0}
|
||||
]
|
||||
|
||||
graph_data = generate_risk_graph(risks_with_controls)
|
||||
self.assertIsInstance(graph_data, str)
|
||||
self.assertTrue(len(graph_data) > 1000)
|
||||
@@ -1,11 +1,18 @@
|
||||
from django.test import TestCase, Client
|
||||
from django.urls import reverse
|
||||
from uuid import uuid4
|
||||
from unittest.mock import patch
|
||||
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl, DocumentTemplate
|
||||
from django.conf import settings
|
||||
from django.contrib.auth.models import User
|
||||
from django.http import HttpResponse
|
||||
|
||||
|
||||
class DocumentViewTest(TestCase):
|
||||
def setUp(self):
|
||||
self.client = Client()
|
||||
self.staff_user = User.objects.create_user(username='staff', password='password', is_staff=True)
|
||||
self.client.login(username='staff', password='password')
|
||||
|
||||
self.organization = Organization.objects.create(
|
||||
id=1,
|
||||
@@ -40,19 +47,43 @@ class DocumentViewTest(TestCase):
|
||||
self.control1 = Control.objects.create(id=1, name="Control A")
|
||||
self.control2 = Control.objects.create(id=2, name="Control B")
|
||||
|
||||
DocumentRiskControl.objects.create(id=1, document=self.document, risk=self.risk1, control=self.control1, weight=5)
|
||||
DocumentRiskControl.objects.create(id=2, document=self.document, risk=self.risk1, control=self.control2, weight=7)
|
||||
DocumentRiskControl.objects.create(id=3, document=self.document, risk=self.risk2, control=self.control1, weight=8)
|
||||
DocumentRiskControl.objects.create(id=1, document=self.document, risk=self.risk1, control=self.control1, weight=5, likelihood=3)
|
||||
DocumentRiskControl.objects.create(id=2, document=self.document, risk=self.risk1, control=self.control2, weight=7, likelihood=4)
|
||||
DocumentRiskControl.objects.create(id=3, document=self.document, risk=self.risk2, control=self.control1, weight=8, likelihood=2)
|
||||
|
||||
template_content = """
|
||||
- segment_type: "h1"
|
||||
content: "{{ document.organization.name }} - Risk Report"
|
||||
- segment_type: "body"
|
||||
content: "Document ID: {{ document.id }}"
|
||||
- segment_type: "body"
|
||||
- segment_type: "p"
|
||||
content: "Created at: {{ document.created_at|date:'Y-m-d' }}"
|
||||
- segment_type: "h2"
|
||||
content: "Risks"
|
||||
content: "Top 10 Risk Identified"
|
||||
- segment_type: "table"
|
||||
content: |
|
||||
<table>
|
||||
<tr>
|
||||
<th>Risk ID</th>
|
||||
<th>Risk Name</th>
|
||||
<th>Inherent Impact</th>
|
||||
<th>Inherent Likelihood</th>
|
||||
<th>Inherent Risk Score</th>
|
||||
<th>Description of Risk</th>
|
||||
</tr>
|
||||
{% for item in risks_with_controls %}
|
||||
<tr>
|
||||
<td>{{ item.risk.id }}</td>
|
||||
<td>{{ item.risk.name }}</td>
|
||||
<td> - </td>
|
||||
<td> - </td>
|
||||
<td> - </td>
|
||||
<td> - </td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</table>
|
||||
- segment_type: "image"
|
||||
content: "data:image/png;base64,{{ graph }}"
|
||||
- segment_type: "h2"
|
||||
content: "Risks with Controls"
|
||||
- segment_type: "body"
|
||||
content: |
|
||||
{% for item in risks_with_controls %}
|
||||
@@ -73,12 +104,44 @@ class DocumentViewTest(TestCase):
|
||||
self.template = DocumentTemplate.objects.create(id=1, name="Default Template", content=template_content)
|
||||
|
||||
def test_document_view(self):
|
||||
url = reverse('core:document', kwargs={'document_id': self.document.id})
|
||||
response = self.client.get(url)
|
||||
response = self.client.get(reverse('core:document', kwargs={'document_id': self.document.id}))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertContains(response, str(self.document.id))
|
||||
self.assertContains(response, self.organization.name)
|
||||
self.assertContains(response, self.document.created_at.strftime('%Y-%m-%d'))
|
||||
self.assertContains(response, self.risk1.risk_name)
|
||||
self.assertContains(response, self.control1.name)
|
||||
self.assertContains(response, "Weight: 5")
|
||||
self.assertTemplateUsed(response, 'document.html')
|
||||
self.assertContains(response, self.organization.name)
|
||||
|
||||
def test_index_view(self):
|
||||
response = self.client.get(reverse('core:index'))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTemplateUsed(response, 'index.html')
|
||||
|
||||
def test_signup_view_get(self):
|
||||
response = self.client.get(reverse('core:signup'))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTemplateUsed(response, 'signup.html')
|
||||
|
||||
def test_thankyou_view(self):
|
||||
response = self.client.get(reverse('core:thankyou'))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTemplateUsed(response, 'thankyou.html')
|
||||
|
||||
def test_payment_page_view_get(self):
|
||||
response = self.client.get(reverse('core:payment_page') + '?email=test@example.com')
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTemplateUsed(response, 'payment.html')
|
||||
|
||||
def test_template_preview_view(self):
|
||||
response = self.client.get(reverse('core:template_preview', args=[self.template.name]))
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTemplateUsed(response, 'template_preview.html')
|
||||
|
||||
def test_pdf_view(self):
|
||||
with patch('backend.core.views.generate_pdf') as mock_generate_pdf:
|
||||
mock_response = HttpResponse(b'%PDF-1.4...', content_type='application/pdf')
|
||||
mock_generate_pdf.return_value = mock_response
|
||||
|
||||
response = self.client.get(reverse('core:pdf_view', args=[self.document.id]))
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response['Content-Type'], 'application/pdf')
|
||||
self.assertIn(b'%PDF', response.content[:10])
|
||||
mock_generate_pdf.assert_called_once_with(self.document)
|
||||
Reference in New Issue
Block a user