56 lines
2.6 KiB
Python
56 lines
2.6 KiB
Python
from unittest.mock import patch
|
|
from django.test import TestCase
|
|
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl
|
|
from backend.accounts.tasks import create_document_for_organization
|
|
|
|
class CeleryTaskTests(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="Category1", primary_impact="High")
|
|
self.control = Control.objects.create(name="Test Control")
|
|
|
|
@patch("backend.accounts.tasks.get_top_risk")
|
|
@patch("backend.accounts.tasks.get_controls_for_risk")
|
|
@patch("backend.accounts.tasks.send_payment_email")
|
|
def test_create_document_for_organization(self, mock_send_payment_email, mock_get_controls_for_risk, mock_get_top_risk):
|
|
mock_get_top_risk.return_value = [self.risk.risk_id]
|
|
mock_get_controls_for_risk.return_value = [(self.control.id, 5, 7)]
|
|
create_document_for_organization(self.organization.email)
|
|
|
|
document = Document.objects.first()
|
|
self.assertIsNotNone(document)
|
|
|
|
document_risk_control = DocumentRiskControl.objects.first()
|
|
self.assertIsNotNone(document_risk_control)
|
|
self.assertEqual(document_risk_control.document, document)
|
|
self.assertEqual(document_risk_control.risk, self.risk)
|
|
self.assertEqual(document_risk_control.control, self.control)
|
|
self.assertEqual(document_risk_control.weight, 5)
|
|
self.assertEqual(document_risk_control.likelihood, 7)
|
|
|
|
mock_send_payment_email.assert_called_once_with(self.organization.email) |