old-riskletpy/backend/core/tests/test_utils.py

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