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

from django.test import TestCase, Client
from django.urls import reverse
from uuid import uuid4
from unittest.mock import patch
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl, DocumentTemplate
from django.conf import settings
from django.contrib.auth.models import User
from django.http import HttpResponse


class DocumentViewTest(TestCase):
    def setUp(self):
        self.client = Client()
        self.staff_user = User.objects.create_user(username='staff', password='password', is_staff=True)
        self.client.login(username='staff', password='password')

        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.document = Document.objects.create(id=uuid4(), organization=self.organization)

        self.risk1 = Risk.objects.create(risk_id=1, risk_name="Risk 1")
        self.risk2 = Risk.objects.create(risk_id=2, risk_name="Risk 2")
        self.control1 = Control.objects.create(id=1, subcategory="PR.AA-01", function="Identity Management")
        self.control2 = Control.objects.create(id=2, subcategory="PR.DS-11", function="Backups")

        DocumentRiskControl.objects.create(id=1, document=self.document, risk=self.risk1, control=self.control1, weight=5, likelihood=3)
        DocumentRiskControl.objects.create(id=2, document=self.document, risk=self.risk1, control=self.control2, weight=7, likelihood=4)
        DocumentRiskControl.objects.create(id=3, document=self.document, risk=self.risk2, control=self.control1, weight=8, likelihood=2)

        template_content = """
        - segment_type: "h1"
          content: 
            title: "{{ document.organization.name }} - Risk Report"
        - segment_type: "p"
          content:
            descripton: |
              "Created at: {{ document.created_at|date:'Y-m-d' }}"
        - segment_type: "h2"
          content: "Top 10 Risk Identified"
        - segment_type: "table"
          html: |
            <table>
              <tr>
                <th>Risk ID</th>
                <th>Risk Name</th>
                <th>Inherent Impact</th>
                <th>Inherent Likelihood</th>
                <th>Inherent Risk Score</th>
                <th>Description of Risk</th>
              </tr>
              {% for item in risks_with_controls %}
              <tr>
                <td>{{ item.risk.id }}</td>
                <td>{{ item.risk.name }}</td>
                <td> - </td>
                <td> - </td>
                <td> - </td>
                <td> - </td>
              </tr>
              {% endfor %}
            </table>
        - segment_type: "image"
          content:
            image: "data:image/png;base64,{{ graph }}"
        - segment_type: "h2"
          content:
            title: "Risks with Controls"
        - segment_type: "body"
          html: |
            {% for item in risks_with_controls %}
            <div class="risk">
              <h3>Risk: {{ item.risk.name }}</h3>
              <div class="controls">
                <h4>Mitigating Controls:</h4>
                {% for control in item.controls %}
                  <div class="control">
                    <span class="name">{{ control.control__subcategory }} - {{ control.control__function }}</span> -
                    <span class="weight">Weight: {{ control.weight }}</span>
                  </div>
                {% endfor %}
              </div>
            </div>
            {% endfor %}
        """
        self.template = DocumentTemplate.objects.create(id=1, name="Default Template", content=template_content)

    def test_document_view(self):
        response = self.client.get(reverse('core:document', kwargs={'document_id': self.document.id}))
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'document.html')
        self.assertContains(response, self.organization.name)  
        self.assertContains(response, "Risk Report")  

    def test_index_view(self):
        response = self.client.get(reverse('core:index'))
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'index.html')

    def test_signup_view_get(self):
        response = self.client.get(reverse('core:signup'))  
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'signup.html')

    def test_thankyou_view(self):
        response = self.client.get(reverse('core:thankyou'))
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'thankyou.html')

    def test_payment_page_view_get(self):
        response = self.client.get(reverse('core:payment_page') + '?email=test@example.com')
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'payment.html')

    def test_template_preview_view(self):
        response = self.client.get(reverse('core:template_preview', args=[self.template.name]))
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'template_preview.html')

    def test_pdf_view(self):
        with patch('backend.core.views.generate_pdf') as mock_generate_pdf:
            mock_response = HttpResponse(b'%PDF-1.4...', content_type='application/pdf')
            mock_generate_pdf.return_value = mock_response

            response = self.client.get(reverse('core:pdf_view', args=[self.document.id]))

            self.assertEqual(response.status_code, 200)
            self.assertEqual(response['Content-Type'], 'application/pdf')
            self.assertIn(b'%PDF', response.content[:10])
            mock_generate_pdf.assert_called_once_with(self.document)