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

class ExportControlsCommandTest(TestCase):
    def setUp(self):
        Control.objects.create(name="Test Safeguard 1")
        Control.objects.create(name="Test Safeguard 2")

        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(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]["CIS v8.1 Safeguards (Sub-Controls)"], "Test Safeguard 1")
            self.assertEqual(rows[1]["CIS v8.1 Safeguards (Sub-Controls)"], "Test Safeguard 2")