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, 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,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__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}, ])