Added graph and table to document

backend/core/tests/test_views.py

@@ -1,11 +1,18 @@
 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,
@@ -40,19 +47,43 @@ class DocumentViewTest(TestCase):
         self.control1 = Control.objects.create(id=1, name="Control A")
         self.control2 = Control.objects.create(id=2, name="Control B")
 
-        DocumentRiskControl.objects.create(id=1, document=self.document, risk=self.risk1, control=self.control1, weight=5)
-        DocumentRiskControl.objects.create(id=2, document=self.document, risk=self.risk1, control=self.control2, weight=7)
-        DocumentRiskControl.objects.create(id=3, document=self.document, risk=self.risk2, control=self.control1, weight=8)
+        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: "{{ document.organization.name }} - Risk Report"
-        - segment_type: "body"
-          content: "Document ID: {{ document.id }}"
-        - segment_type: "body"
+        - segment_type: "p"
           content: "Created at: {{ document.created_at|date:'Y-m-d' }}"
         - segment_type: "h2"
-          content: "Risks"
+          content: "Top 10 Risk Identified"
+        - segment_type: "table"
+          content: |
+            <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: "data:image/png;base64,{{ graph }}"
+        - segment_type: "h2"
+          content: "Risks with Controls"
         - segment_type: "body"
           content: |
             {% for item in risks_with_controls %}
@@ -73,12 +104,44 @@ class DocumentViewTest(TestCase):
         self.template = DocumentTemplate.objects.create(id=1, name="Default Template", content=template_content)
 
     def test_document_view(self):
-        url = reverse('core:document', kwargs={'document_id': self.document.id})
-        response = self.client.get(url)
+        response = self.client.get(reverse('core:document', kwargs={'document_id': self.document.id}))
         self.assertEqual(response.status_code, 200)
-        self.assertContains(response, str(self.document.id))
-        self.assertContains(response, self.organization.name)
-        self.assertContains(response, self.document.created_at.strftime('%Y-%m-%d'))
-        self.assertContains(response, self.risk1.risk_name)
-        self.assertContains(response, self.control1.name)
-        self.assertContains(response, "Weight: 5")
+        self.assertTemplateUsed(response, 'document.html')
+        self.assertContains(response, self.organization.name) 
+
+    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)