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, 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"], sensitive_data_types={ "personal": {"applicable": True, "impact": 4}, "financial": {"applicable": True, "impact": 3}, "ip": {"applicable": False, "impact": None}, "operational": {"applicable": True, "impact": 5}, "government": {"applicable": False, "impact": None}, "none": {"applicable": False} }, 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, subcategory=f"C-{i}", function=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. **Risk ID 1 (Privacy Regulation Violation)**: Critical because the company's operations are governed by NIS2 regulations, and any data breach could lead to severe financial penalties and reputational damage.\n" "2. **Risk ID 2 (Third Party Code Compromise)**: This risk is critical given the company's reliance on more than five third-party vendors, which increases the potential for system compromises and data breaches through external partnerships.\n" "3. **Risk ID 3 (Misconfigured Cloud Services)**: Critical due to the company's hybrid IT infrastructure, which may lead to increased data exposure if cloud services are not properly configured, impacting compliance and customer trust.\n" ) mock_client.chat.completions.create.return_value = mock_response risks = get_top_risk(self.organization) top_risk_ids = [r['risk_id'] for r in risks] self.assertEqual(top_risk_ids, [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__subcategory': 'PR.AA-01', 'control__function': 'Identity'}, {'control': 102, 'control__subcategory': 'PR.DS-11', 'control__function': 'Backups'}, ] }, { 'risk': {'id': 2, 'name': 'Risk 2'}, 'controls': [ {'control': 101, 'control__subcategory': 'PR.AA-01', 'control__function': 'Identity'}, ] } ] summary = get_safeguard_summary_table(risks_with_controls) self.assertEqual(summary, [ {'id': 101, 'subcategory': '', 'category': '', 'function': '', 'name': 'PR.AA-01 - Identity', 'count': 2}, {'id': 102, 'subcategory': '', 'category': '', 'function': '', 'name': 'PR.DS-11 - Backups', 'count': 1}, ])