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)