fix query-set-extra tests (#3635)

Co-authored-by: Lewis <[email protected]>

Author: tean-lai

python/django/security/audit/query-set-extra.py

@@ -14,34 +14,74 @@
 Entry.objects.get({}).filter().update().extra()
 
 # ok:avoid-query-set-extra
-findings = Finding.objects.filter(verified=True,
-                                      severity__in=('Critical', 'High', 'Medium', 'Low', 'Info')).prefetch_related(
-        'test__engagement__product',
-        'test__engagement__product__prod_type',
-        'test__engagement__risk_acceptance',
-        'risk_acceptance_set',
-        'reporter').extra(
+findings = (
+    Finding.objects.filter(
+        verified=True, severity__in=("Critical", "High", "Medium", "Low", "Info")
+    )
+    .prefetch_related(
+        "test__engagement__product",
+        "test__engagement__product__prod_type",
+        "test__engagement__risk_acceptance",
+        "risk_acceptance_set",
+        "reporter",
+    )
+    .extra(
         select={
-            'ra_count': 'SELECT COUNT(*) FROM dojo_risk_acceptance INNER JOIN '
-                        'dojo_risk_acceptance_accepted_findings ON '
-                        '( dojo_risk_acceptance.id = dojo_risk_acceptance_accepted_findings.risk_acceptance_id ) '
-                        'WHERE dojo_risk_acceptance_accepted_findings.finding_id = dojo_finding.id',
+            "ra_count": "SELECT COUNT(*) FROM dojo_risk_acceptance INNER JOIN "
+            "dojo_risk_acceptance_accepted_findings ON "
+            "( dojo_risk_acceptance.id = dojo_risk_acceptance_accepted_findings.risk_acceptance_id ) "
+            "WHERE dojo_risk_acceptance_accepted_findings.finding_id = dojo_finding.id",
         },
     )
+)
 
-example = 1
+
+example = input()
 # ruleid:avoid-query-set-extra
-active_findings = Finding.objects.filter(verified=True, active=True,
-                                      severity__in=('Critical', 'High', 'Medium', 'Low', 'Info')).prefetch_related(
-        'test__engagement__product',
-        'test__engagement__product__prod_type',
-        'test__engagement__risk_acceptance',
-        'risk_acceptance_set',
-        'reporter').extra(
+active_findings = (
+    Finding.objects.filter(
+        verified=True,
+        active=True,
+        severity__in=("Critical", "High", "Medium", "Low", "Info"),
+    )
+    .prefetch_related(
+        "test__engagement__product",
+        "test__engagement__product__prod_type",
+        "test__engagement__risk_acceptance",
+        "risk_acceptance_set",
+        "reporter",
+    )
+    .extra(
+        select={
+            "ra_count": f"SELECT COUNT(*) FROM dojo_risk_acceptance INNER JOIN "
+            f"dojo_risk_acceptance_accepted_findings ON "
+            f"( dojo_risk_acceptance.id = dojo_risk_acceptance_accepted_findings.risk_acceptance_id ) "
+            f"WHERE dojo_risk_acceptance_accepted_findings.finding_id = {example}",
+        },
+    )
+)
+
+example = 1
+# ok:avoid-query-set-extra
+active_findings = (
+    Finding.objects.filter(
+        verified=True,
+        active=True,
+        severity__in=("Critical", "High", "Medium", "Low", "Info"),
+    )
+    .prefetch_related(
+        "test__engagement__product",
+        "test__engagement__product__prod_type",
+        "test__engagement__risk_acceptance",
+        "risk_acceptance_set",
+        "reporter",
+    )
+    .extra(
         select={
-            'ra_count': f'SELECT COUNT(*) FROM dojo_risk_acceptance INNER JOIN '
-                        f'dojo_risk_acceptance_accepted_findings ON '
-                        f'( dojo_risk_acceptance.id = dojo_risk_acceptance_accepted_findings.risk_acceptance_id ) '
-                        f'WHERE dojo_risk_acceptance_accepted_findings.finding_id = {example}',
+            "ra_count": f"SELECT COUNT(*) FROM dojo_risk_acceptance INNER JOIN "
+            f"dojo_risk_acceptance_accepted_findings ON "
+            f"( dojo_risk_acceptance.id = dojo_risk_acceptance_accepted_findings.risk_acceptance_id ) "
+            f"WHERE dojo_risk_acceptance_accepted_findings.finding_id = {example}",
         },
     )
+)