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, 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.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: | {% for item in risks_with_controls %} {% endfor %}
Risk ID Risk Name Inherent Impact Inherent Likelihood Inherent Risk Score Description of Risk
{{ item.risk.id }} {{ item.risk.name }} - - - -
- 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 %}

Risk: {{ item.risk.name }}

Mitigating Controls:

{% for control in item.controls %}
{{ control.control__subcategory }} - {{ control.control__function }} - Weight: {{ control.weight }}
{% endfor %}
{% 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)