import os
import csv
from django.core.management import call_command
from django.test import TestCase
from backend.core.models import Risk, Control

class ExportControlsCommandTest(TestCase):
    def setUp(self):
        self.risk1 = Risk.objects.create(risk_id=1, risk_name="Test Risk 1")
        self.risk2 = Risk.objects.create(risk_id=2, risk_name="Test Risk 2")
        
        self.control1 = Control.objects.create(
            risk=self.risk1,
            safeguard="Test Safeguard 1",
            description="Description 1",
            weight=5
        )
        self.control2 = Control.objects.create(
            risk=self.risk2,
            safeguard="Test Safeguard 2",
            description="Description 2",
            weight=10
        )

        self.csv_file_path = 'test_export_controls.csv'

    def tearDown(self):
        if os.path.exists(self.csv_file_path):
            os.remove(self.csv_file_path)

    def test_export_controls_success(self):
        call_command('export_controls', self.csv_file_path)

        self.assertTrue(os.path.exists(self.csv_file_path))

        with open(self.csv_file_path, mode='r', encoding='utf-8') as csv_file:
            reader = csv.DictReader(csv_file)
            rows = list(reader)

            self.assertEqual(len(rows), 2) 

            self.assertEqual(rows[0]["Risk #"], str(self.risk1.risk_id))
            self.assertEqual(rows[0]["Risk Description"], self.risk1.risk_name)
            self.assertEqual(rows[0]["CIS v8.1 Safeguards (Sub-Controls)"], self.control1.safeguard)
            self.assertEqual(rows[0]["Weight (0-10)"], str(self.control1.weight))

            self.assertEqual(rows[1]["Risk #"], str(self.risk2.risk_id))
            self.assertEqual(rows[1]["Risk Description"], self.risk2.risk_name)
            self.assertEqual(rows[1]["CIS v8.1 Safeguards (Sub-Controls)"], self.control2.safeguard)
            self.assertEqual(rows[1]["Weight (0-10)"], str(self.control2.weight))

    def test_export_controls_empty(self):
        Control.objects.all().delete()

        call_command('export_controls', self.csv_file_path)

        self.assertTrue(os.path.exists(self.csv_file_path))

        with open(self.csv_file_path, mode='r', encoding='utf-8') as csv_file:
            reader = csv.DictReader(csv_file)
            rows = list(reader)

            self.assertEqual(len(rows), 0) 