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

class ImportControlsCommandTest(TestCase):
    def setUp(self):
        self.csv_file_path = 'test_import_controls.csv'
        with open(self.csv_file_path, mode='w', encoding='utf-8') as csv_file:
            writer = csv.DictWriter(csv_file, fieldnames=["CIS v8.1 Safeguards (Sub-Controls)"])
            writer.writeheader()
            writer.writerow({"CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 1"})
            writer.writerow({"CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 2"})

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

    def test_import_controls(self):
        self.assertEqual(Control.objects.count(), 0)

        call_command('import_controls', self.csv_file_path)

        self.assertEqual(Control.objects.count(), 2)

        safeguards = Control.objects.values_list('name', flat=True)
        self.assertIn("Test Safeguard 1", safeguards)
        self.assertIn("Test Safeguard 2", safeguards)

    def test_import_controls_update(self):
        Control.objects.create(name="Test Safeguard 1")

        call_command('import_controls', self.csv_file_path)

        self.assertEqual(Control.objects.count(), 2)

        safeguards = Control.objects.values_list('name', flat=True)
        self.assertIn("Test Safeguard 1", safeguards)
        self.assertIn("Test Safeguard 2", safeguards)