172 lines
6.4 KiB
Python
172 lines
6.4 KiB
Python
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} : 4 : 3" 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, 50)
|
|
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': 4.0, 'likelihood': 4.0}
|
|
]
|
|
|
|
graph_data = generate_risk_graph(risks_with_controls)
|
|
self.assertIsInstance(graph_data, str)
|
|
self.assertTrue(len(graph_data) > 1000)
|
|
|
|
def test_generate_residual_risk_graph_base64(self):
|
|
risks_with_controls = [
|
|
{
|
|
'risk': {'id': 1, 'name': 'Risk 1'},
|
|
'residual_impact': 3,
|
|
'residual_likelihood': 4,
|
|
},
|
|
{
|
|
'risk': {'id': 2, 'name': 'Risk 2'},
|
|
'residual_impact': 2,
|
|
'residual_likelihood': 2,
|
|
}
|
|
]
|
|
|
|
graph_data = generate_residual_risk_graph(risks_with_controls)
|
|
self.assertIsInstance(graph_data, str)
|
|
self.assertTrue(len(graph_data) > 1000)
|
|
|
|
def test_get_safeguard_summary_table_basic(self):
|
|
from backend.core.tables import get_safeguard_summary_table
|
|
risks_with_controls = [
|
|
{
|
|
'risk': {'id': 1, 'name': 'Risk 1'},
|
|
'controls': [
|
|
{'control': 101, 'control__name': 'Control A'},
|
|
{'control': 102, 'control__name': 'Control B'},
|
|
]
|
|
},
|
|
{
|
|
'risk': {'id': 2, 'name': 'Risk 2'},
|
|
'controls': [
|
|
{'control': 101, 'control__name': 'Control A'},
|
|
]
|
|
}
|
|
]
|
|
summary = get_safeguard_summary_table(risks_with_controls)
|
|
self.assertEqual(summary, [
|
|
{'id': 101, 'name': 'Control A', 'count': 2},
|
|
{'id': 102, 'name': 'Control B', 'count': 1},
|
|
]) |