Added graph and table to document

backend/core/tests/test_utils.py

@@ -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)