Fix typo in swap function and fix migration test

Signed-off-by: ziadhany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/migrations/0030_vulnerabilityseverity_rm_extra_records_swap.py

@@ -1,4 +1,5 @@
 from django.db import migrations
+from vulnerabilities.severity_systems import CVSSV2_VECTOR, CVSSV31_VECTOR, CVSSV3, CVSSV2, CVSSV31, CVSSV3_VECTOR
 
 
 def remove_extra_rows(apps, schema_editor):
@@ -21,12 +22,17 @@ def remove_extra_rows(apps, schema_editor):
 def swap_scoring_elements_with_value(apps, schema_editor):
     VulnerabilitySeverity = apps.get_model('vulnerabilities', 'VulnerabilitySeverity')
     for vuln_severity in VulnerabilitySeverity.objects.all():
+
         cvss_mapper = {
-            "cvssv2_vector": "cvssv2",
-            "cvssv3_vector": "cvssv3",
-            "cvssv3.1_vector": "cvssv3.1",
+            CVSSV2_VECTOR.identifier: CVSSV2.identifier,
+            CVSSV3_VECTOR.identifier: CVSSV3.identifier,
+            CVSSV31_VECTOR.identifier: CVSSV31.identifier,
         }
-        if VulnerabilitySeverity.scoring_system in ["cvssv2_vector", "cvssv3_vector", "cvssv3.1_vector"]:
+
+        if vuln_severity.scoring_system in [CVSSV2_VECTOR.identifier,
+                                            CVSSV3_VECTOR.identifier,
+                                            CVSSV31_VECTOR.identifier]:
+
             vuln_severity.scoring_system = cvss_mapper[vuln_severity.scoring_system]
             temp = vuln_severity.scoring_elements  # value
             vuln_severity.scoring_elements = vuln_severity.value

vulnerabilities/tests/test_data_migrations.py

@@ -275,27 +275,27 @@ def test_rm_extra_records_and_swap(self):
         # using get_model to avoid circular import
         VulnerabilitySeverity = self.apps.get_model("vulnerabilities", "VulnerabilitySeverity")
         severity_cvss2 = VulnerabilitySeverity.objects.get(reference_id=1)
-        self.assertEqual(severity_cvss2.value, "AV:N/AC:L/Au:N/C:P/I:P/A:P")
-        self.assertEqual(severity_cvss2.scoring_elements, "7.5")
+        self.assertEqual(severity_cvss2.scoring_elements, "AV:N/AC:L/Au:N/C:P/I:P/A:P")
+        self.assertEqual(severity_cvss2.value, "7.5")
 
         severity_cvss3 = VulnerabilitySeverity.objects.get(reference_id=2)
-        self.assertEqual(severity_cvss3.value, "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H")
-        self.assertEqual(severity_cvss3.scoring_elements, "7.5")
+        self.assertEqual(severity_cvss3.scoring_elements, "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H")
+        self.assertEqual(severity_cvss3.value, "7.5")
 
         severity_cvss3_1 = VulnerabilitySeverity.objects.get(reference_id=3)
-        self.assertEqual(severity_cvss3_1.value, "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H")
-        self.assertEqual(severity_cvss3_1.scoring_elements, "9.8")
+        self.assertEqual(severity_cvss3_1.scoring_elements, "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H")
+        self.assertEqual(severity_cvss3_1.value, "9.8")
 
         severities = VulnerabilitySeverity.objects.filter(reference_id=4)
         assert len(severities) == 5
-        self.assertEqual(severities[0].value, "AV:N/AC:L/Au:N/C:P/I:P/A:P")
-        self.assertEqual(severities[0].scoring_elements, "7.5")
+        self.assertEqual(severities[0].scoring_elements, "AV:N/AC:L/Au:N/C:P/I:P/A:P")
+        self.assertEqual(severities[0].value, "7.5")
 
-        self.assertEqual(severities[1].value, "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H")
-        self.assertEqual(severities[1].scoring_elements, "7.5")
+        self.assertEqual(severities[1].scoring_elements, "CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H")
+        self.assertEqual(severities[1].value, "7.5")
 
-        self.assertEqual(severities[2].value, "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H")
-        self.assertEqual(severities[2].scoring_elements, "9.8")
+        self.assertEqual(severities[2].scoring_elements, "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H")
+        self.assertEqual(severities[2].value, "9.8")
 
         self.assertEqual(severities[3].value, "Low")
         self.assertIsNone(severities[3].scoring_elements)