from django.test import TestCase

from utilities.testing import ChangeLoggedFilterSetTests

from nb_risk.models import Vulnerability
from nb_risk.filters import VulnerabilityFilterSet
from nb_risk import choices


class VulnerabilityFilterTestCase(TestCase, ChangeLoggedFilterSetTests):
    queryset = Vulnerability.objects.all()
    filterset = VulnerabilityFilterSet

    @classmethod
    def setUpTestData(cls):

        vulnerabilities = (
            Vulnerability(
                name="Vulnerability 1",
                cve="CVE-2021-1234",
                description="A test vulnerability",
            ),
            Vulnerability(
                name="Vulnerability 2",
                cve="CVE-2021-1235",
                description="Another test vulnerability",
            ),
            Vulnerability(
                name="Vulnerability 3",
                cve="CVE-2021-1236",
                description="Yet another test vulnerability",
            ),
        )
        Vulnerability.objects.bulk_create(vulnerabilities)

    def test_name_none(self):
        params = {"name": ["None"]}
        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)

    def test_name(self):
        params = {"name": ["Vulnerability 2", "Vulnerability 3"]}
        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 2)

    def test_names(self):
        params = {"name__icontains": ["Vulnerability"]}
        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 3)

    def test_cve(self):
        params = {"cve": ["CVE-2021-1236"]}
        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 1)

    def test_cve_none(self):
        params = {"cve": ["12345"]}
        self.assertEqual(self.filterset(params, self.queryset).qs.count(), 0)
