From 988d99bdb8c3489ccd172146601f4692d87cbc91 Mon Sep 17 00:00:00 2001
From: Valentijn Scholten <valentijnscholten@gmail.com>
Date: Fri, 27 Feb 2026 11:19:22 +0100
Subject: [PATCH 1/5] Add v3 Locations performance test and update script
 support

- Add TestDojoImporterPerformanceSmallLocations with V3_FEATURE_LOCATIONS
- Update update_performance_test_counts.py to run both v2 and v3 test classes
- Add --no-keepdb and EXTRA_ARGS to run-unittest.sh for test flexibility
---
 run-unittest.sh                           |  10 +-
 scripts/update_performance_test_counts.py | 215 +++++++++++++---------
 unittests/test_importers_performance.py   | 145 ++++++++++++++-
 3 files changed, 277 insertions(+), 93 deletions(-)

diff --git a/run-unittest.sh b/run-unittest.sh
index 6073d4ff582..9ffeeeb20b0 100755
--- a/run-unittest.sh
+++ b/run-unittest.sh
@@ -1,6 +1,8 @@
 #!/usr/bin/env bash
 unset TEST_CASE
 unset FAIL_FAST
+unset KEEP_DB
+EXTRA_ARGS=()
 
 bash ./docker/docker-compose-check.sh
 if [[ $? -eq 1 ]]; then exit 1; fi
@@ -12,6 +14,7 @@ usage() {
   echo "Options:"
   echo "  --test-case -t {YOUR_FULLY_QUALIFIED_TEST_CASE}"
   echo "  --fail-fast -f - stop on first test failure"
+  echo "  --no-keepdb - recreate the test database (don't reuse existing)"
   echo "  --help -h - prints this dialogue."
   echo
   echo "You must specify a test case (arg)!"
@@ -41,6 +44,10 @@ while [[ $# -gt 0 ]]; do
       FAIL_FAST="--failfast"
       shift # past argument
       ;;
+    --no-keepdb)
+      KEEP_DB=""
+      shift # past argument
+      ;;
     -h|--help)
       usage
       exit 0
@@ -66,4 +73,5 @@ echo "Running docker compose unit tests with test case $TEST_CASE ..."
 # Compose V2 integrates compose functions into the Docker platform, continuing to support
 # most of the  previous docker-compose features and flags. You can run Compose V2 by
 # replacing the hyphen (-) with a space, using docker compose, instead of docker-compose.
-docker compose exec uwsgi bash -c "python manage.py test $TEST_CASE -v2 --keepdb $FAIL_FAST"
+KEEP_DB="${KEEP_DB:---keepdb}"
+docker compose exec uwsgi bash -c "python manage.py test $TEST_CASE -v2 $KEEP_DB $FAIL_FAST ${EXTRA_ARGS[*]}"
diff --git a/scripts/update_performance_test_counts.py b/scripts/update_performance_test_counts.py
index bc99e108da8..9b406539e6e 100644
--- a/scripts/update_performance_test_counts.py
+++ b/scripts/update_performance_test_counts.py
@@ -11,20 +11,18 @@
 
 How to run:
 
-    # Default: Update the test file (uses TestDojoImporterPerformanceSmall by default)
+    # Default: Update both v2 and v3 test classes
     python3 scripts/update_performance_test_counts.py
 
-    # Or specify a different test class:
-    python3 scripts/update_performance_test_counts.py --test-class TestDojoImporterPerformanceSmall
-
     # Step 1: Run tests and generate report only (without updating)
     python3 scripts/update_performance_test_counts.py --report-only
 
     # Step 2: Verify all tests pass
     python3 scripts/update_performance_test_counts.py --verify
 
-The script defaults to TestDojoImporterPerformanceSmall if --test-class is not provided.
-The script defaults to --update behavior if no action flag is provided.
+The script always runs and updates both TestDojoImporterPerformanceSmall (v2) and
+TestDojoImporterPerformanceSmallLocations (v3). The script defaults to --update
+behavior if no action flag is provided.
 """
 
 import argparse
@@ -36,6 +34,12 @@
 # Path to the test file
 TEST_FILE = Path(__file__).parent.parent / "unittests" / "test_importers_performance.py"
 
+# Both v2 and v3 performance test classes - script always updates/verifies both
+TEST_CLASSES = (
+    "TestDojoImporterPerformanceSmall",
+    "TestDojoImporterPerformanceSmallLocations",
+)
+
 
 class TestCount:
 
@@ -64,9 +68,9 @@ def extract_test_methods(test_class: str) -> list[str]:
 
     content = TEST_FILE.read_text()
 
-    # Find the test class definition
+    # Find the test class definition (use (?<!-) to avoid matching "test-class" in docstrings)
     class_pattern = re.compile(
-        rf"class {re.escape(test_class)}.*?(?=class |\Z)",
+        rf"class {re.escape(test_class)}.*?(?=(?<!-)class \w|\Z)",
         re.DOTALL,
     )
     class_match = class_pattern.search(content)
@@ -362,12 +366,17 @@ def generate_report(counts: list[TestCount], expected_counts: dict[str, dict[str
 
     print("=" * 80)
     print("\nTo update the test file, run:")
-    print(f"  python scripts/update_performance_test_counts.py --test-class {test_name.split('_')[0]} --update")
+    print("  python3 scripts/update_performance_test_counts.py")
     print()
 
 
-def update_test_file(counts: list[TestCount]):
-    """Update the test file with new expected counts."""
+def update_test_file(counts: list[TestCount], test_class: str | None = None):
+    """
+    Update the test file with new expected counts.
+
+    When test_class is provided, only update test methods within that class.
+    This is required when multiple classes (e.g. v2 and v3) share the same test method names.
+    """
     if not counts:
         print("No counts to update.")
         return
@@ -419,22 +428,37 @@ def _extract_call_span(method_content: str, call_name: str) -> tuple[int, int] |
         "second_import_async_tasks": "expected_num_async_tasks2",
     }
 
+    # Restrict search to the specified test class if given
+    search_content = content
+    search_offset = 0
+    if test_class:
+        class_pattern = re.compile(
+            rf"class {re.escape(test_class)}.*?(?=(?<!-)class \w|\Z)",
+            re.DOTALL,
+        )
+        class_match = class_pattern.search(content)
+        if not class_match:
+            print(f"⚠️  Warning: Could not find test class {test_class}")
+            return
+        search_content = class_match.group(0)
+        search_offset = class_match.start()
+
     # Update each test method
     for test_name, test_updates in updates.items():
         print(f"  Updating {test_name}...")
-        # Find the test method boundaries
+        # Find the test method boundaries within the search scope
         test_method_pattern = re.compile(
             rf"(def {re.escape(test_name)}\([^)]*\):.*?)(?=def test_|\Z)",
             re.DOTALL,
         )
-        test_match = test_method_pattern.search(content)
+        test_match = test_method_pattern.search(search_content)
         if not test_match:
             print(f"⚠️  Warning: Could not find test method {test_name}")
             continue
 
         test_method_content = test_match.group(1)
-        test_method_start = test_match.start()
-        test_method_end = test_match.end()
+        test_method_start = search_offset + test_match.start()
+        test_method_end = search_offset + test_match.end()
 
         call_span = _extract_call_span(test_method_content, "self._import_reimport_performance")
         param_map = param_map_import_reimport
@@ -543,12 +567,6 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog=__doc__,
     )
-    parser.add_argument(
-        "--test-class",
-        required=False,
-        default="TestDojoImporterPerformanceSmall",
-        help="Test class name (e.g., TestDojoImporterPerformanceSmall). Defaults to TestDojoImporterPerformanceSmall if not provided.",
-    )
     parser.add_argument(
         "--report-only",
         action="store_true",
@@ -568,81 +586,88 @@ def main():
     args = parser.parse_args()
 
     if args.report_only:
-        # Step 1: Run tests and generate report
-        # Run each test method individually
-        test_methods = extract_test_methods(args.test_class)
-        if not test_methods:
-            print(f"⚠️  No test methods found in {args.test_class}")
-            sys.exit(1)
-
-        print(f"\nFound {len(test_methods)} test method(s) in {args.test_class}")
-        print("=" * 80)
-
+        # Step 1: Run tests and generate report for both test classes
         all_counts = []
-        for test_method in test_methods:
-            print(f"\n{'=' * 80}")
-            output, return_code = run_test_method(args.test_class, test_method)
-            success, error_msg = check_test_execution_success(output, return_code)
-            if not success:
-                print(f"\n⚠️  Test execution failed for {test_method}: {error_msg}")
-                print("Skipping this test method...")
+        for test_class in TEST_CLASSES:
+            test_methods = extract_test_methods(test_class)
+            if not test_methods:
+                print(f"⚠️  No test methods found in {test_class}")
                 continue
 
-            counts = parse_test_output(output)
-            if counts:
-                all_counts.extend(counts)
+            print(f"\n{'=' * 80}")
+            print(f"Test class: {test_class} ({len(test_methods)} methods)")
+            print("=" * 80)
+
+            for test_method in test_methods:
+                print("\n---")
+                output, return_code = run_test_method(test_class, test_method)
+                success, error_msg = check_test_execution_success(output, return_code)
+                if not success:
+                    print(f"\n⚠️  Test execution failed for {test_method}: {error_msg}")
+                    print("Skipping this test method...")
+                    continue
+
+                counts = parse_test_output(output)
+                if counts:
+                    all_counts.extend(counts)
 
-        expected_counts = extract_expected_counts_from_file(args.test_class)
-        generate_report(all_counts, expected_counts)
+        if all_counts:
+            generate_report(all_counts, {})
+        else:
+            print("✅ All tests passed! No count differences found.")
 
     elif args.verify:
-        # Step 3: Verify
-        success = verify_tests(args.test_class)
-        sys.exit(0 if success else 1)
+        # Step 3: Verify both test classes
+        all_pass = True
+        for test_class in TEST_CLASSES:
+            if not verify_tests(test_class):
+                all_pass = False
+        sys.exit(0 if all_pass else 1)
 
     else:
         # Default: Update the file (--update is the default behavior)
-        # Run each test method individually
-        test_methods = extract_test_methods(args.test_class)
-        if not test_methods:
-            print(f"⚠️  No test methods found in {args.test_class}")
-            sys.exit(1)
-
-        print(f"\nFound {len(test_methods)} test method(s) in {args.test_class}")
-        print("=" * 80)
-
+        # Run each test method in both test classes
         all_counts = []
-        for test_method in test_methods:
-            print(f"\n{'=' * 80}")
-            output, return_code = run_test_method(args.test_class, test_method)
-            success, error_msg = check_test_execution_success(output, return_code)
-            if not success:
-                print(f"\n⚠️  Test execution failed for {test_method}: {error_msg}")
-                print("Skipping this test method...")
+        for test_class in TEST_CLASSES:
+            test_methods = extract_test_methods(test_class)
+            if not test_methods:
+                print(f"⚠️  No test methods found in {test_class}")
                 continue
 
-            counts = parse_test_output(output)
-
-            # Check if test actually passed
-            test_passed = "OK" in output or ("Ran" in output and "FAILED" not in output and return_code == 0)
-
-            if counts:
-                all_counts.extend(counts)
-                # Update immediately after each test
-                update_test_file(counts)
-                print(f"⚠️  {test_method}: Found {len(counts)} count mismatch(es) - updated file")
-            elif test_passed:
-                print(f"✅ {test_method}: Test passed, all counts match")
-            elif return_code != 0:
-                # Test might have failed for other reasons
-                print(f"⚠️  {test_method}: Test failed (exit code {return_code}) but no count mismatches parsed")
-                print("   This might indicate a parsing issue or a different type of failure")
-                # Show a snippet of the output to help debug
-                fail_lines = [line for line in output.split("\n") if "FAIL" in line or "Error" in line or "Exception" in line]
-                if fail_lines:
-                    print("   Relevant error lines:")
-                    for line in fail_lines[:5]:
-                        print(f"     {line}")
+            print(f"\n{'=' * 80}")
+            print(f"Test class: {test_class} ({len(test_methods)} methods)")
+            print("=" * 80)
+
+            for test_method in test_methods:
+                print("\n---")
+                output, return_code = run_test_method(test_class, test_method)
+                success, error_msg = check_test_execution_success(output, return_code)
+                if not success:
+                    print(f"\n⚠️  Test execution failed for {test_method}: {error_msg}")
+                    print("Skipping this test method...")
+                    continue
+
+                counts = parse_test_output(output)
+
+                # Check if test actually passed
+                test_passed = "OK" in output or ("Ran" in output and "FAILED" not in output and return_code == 0)
+
+                if counts:
+                    all_counts.extend(counts)
+                    # Update immediately after each test
+                    update_test_file(counts, test_class=test_class)
+                    print(f"⚠️  {test_method}: Found {len(counts)} count mismatch(es) - updated file")
+                elif test_passed:
+                    print(f"✅ {test_method}: Test passed, all counts match")
+                elif return_code != 0:
+                    # Test might have failed for other reasons
+                    print(f"⚠️  {test_method}: Test failed (exit code {return_code}) but no count mismatches parsed")
+                    print("   This might indicate a parsing issue or a different type of failure")
+                    fail_lines = [line for line in output.split("\n") if "FAIL" in line or "Error" in line or "Exception" in line]
+                    if fail_lines:
+                        print("   Relevant error lines:")
+                        for line in fail_lines[:5]:
+                            print(f"     {line}")
 
         if all_counts:
             print(f"\n{'=' * 80}")
@@ -650,14 +675,22 @@ def main():
             # Some performance counts can vary depending on test ordering / keepdb state.
             # Do a final full-suite pass and apply any remaining mismatches so the suite passes as run in CI.
             print("\nRunning a final verify pass for stability...")
-            success, suite_mismatches = verify_and_get_mismatches(args.test_class)
-            if not success and suite_mismatches:
-                print("\nApplying remaining mismatches from full-suite run...")
-                update_test_file(suite_mismatches)
+            all_pass = True
+            for test_class in TEST_CLASSES:
+                success, suite_mismatches = verify_and_get_mismatches(test_class)
+                if not success and suite_mismatches:
+                    print(f"\nApplying remaining mismatches from {test_class}...")
+                    update_test_file(suite_mismatches, test_class=test_class)
+                    all_pass = False
+            if not all_pass:
                 print("\nRe-running verify...")
-                success, _ = verify_and_get_mismatches(args.test_class)
-                sys.exit(0 if success else 1)
-            sys.exit(0 if success else 1)
+                all_pass = True
+                for test_class in TEST_CLASSES:
+                    success, _ = verify_and_get_mismatches(test_class)
+                    if not success:
+                        all_pass = False
+                sys.exit(0 if all_pass else 1)
+            sys.exit(0 if all_pass else 1)
         else:
             print(f"\n{'=' * 80}")
             print("\n✅ No differences found. All tests are already up to date.")
diff --git a/unittests/test_importers_performance.py b/unittests/test_importers_performance.py
index d0bd84101f5..f8da76d1433 100644
--- a/unittests/test_importers_performance.py
+++ b/unittests/test_importers_performance.py
@@ -30,6 +30,7 @@
 from dojo.decorators import dojo_async_task_counter
 from dojo.importers.default_importer import DefaultImporter
 from dojo.importers.default_reimporter import DefaultReImporter
+from dojo.location.models import Location, LocationFindingReference
 from dojo.models import (
     Development_Environment,
     Dojo_User,
@@ -228,7 +229,6 @@ def _import_reimport_performance(
                             test, _, _len_new_findings, _len_closed_findings, _, _, _ = reimporter.process_scan(scan)
 
 
-# TODO: Implement Locations
 @skip_unless_v2
 class TestDojoImporterPerformanceSmall(TestDojoImporterPerformanceBase):
 
@@ -462,3 +462,146 @@ def test_deduplication_performance_pghistory_no_async(self):
             expected_num_queries2=236,
             expected_num_async_tasks2=7,
         )
+
+
+@override_settings(V3_FEATURE_LOCATIONS=True)
+class TestDojoImporterPerformanceSmallLocations(TestDojoImporterPerformanceBase):
+
+    r"""
+    Performance tests using small sample files (StackHawk, ~6 findings) with v3 locations enabled.
+
+    These mirror TestDojoImporterPerformanceSmall but run with V3_FEATURE_LOCATIONS=True.
+    Query counts are specific to the locations code path and will differ from the v2 endpoint counts.
+
+    To determine or update the expected counts, run:
+        python3 scripts/update_performance_test_counts.py --test-class TestDojoImporterPerformanceSmallLocations
+    """
+
+    def setUp(self):
+        super().setUp()
+        # Warm up ContentType cache for Location models so these queries don't
+        # count against the measured sections.
+        for model in [Location, LocationFindingReference]:
+            ContentType.objects.get_for_model(model)
+
+    def _import_reimport_performance(self, expected_num_queries1, expected_num_async_tasks1, expected_num_queries2, expected_num_async_tasks2, expected_num_queries3, expected_num_async_tasks3):
+        r"""
+        Log output can be quite large as when the assertNumQueries fails, all queries are printed.
+        It could be useful to capture the output in `less`:
+            ./run-unittest.sh --test-case unittests.test_importers_performance.TestDojoImporterPerformanceSmallLocations 2>&1 | less
+        Then search for `expected` to find the lines where the expected number of queries is printed.
+        Or you can use `grep` to filter the output:
+            ./run-unittest.sh --test-case unittests.test_importers_performance.TestDojoImporterPerformanceSmallLocations 2>&1 | grep expected -B 10
+        """
+        return super()._import_reimport_performance(
+            expected_num_queries1,
+            expected_num_async_tasks1,
+            expected_num_queries2,
+            expected_num_async_tasks2,
+            expected_num_queries3,
+            expected_num_async_tasks3,
+            scan_file1=STACK_HAWK_SUBSET_FILENAME,
+            scan_file2=STACK_HAWK_FILENAME,
+            scan_file3=STACK_HAWK_SUBSET_FILENAME,
+            scan_type=STACK_HAWK_SCAN_TYPE,
+            product_name="TestDojoDefaultImporterLocations",
+            engagement_name="Test Create Engagement Locations",
+        )
+
+    @override_settings(ENABLE_AUDITLOG=True)
+    def test_import_reimport_reimport_performance_pghistory_async(self):
+        """
+        This test checks the performance of the importers when using django-pghistory with async enabled.
+        Query counts will need to be determined by running the test initially.
+        """
+        configure_audit_system()
+        configure_pghistory_triggers()
+
+        self._import_reimport_performance(
+            expected_num_queries1=1091,
+            expected_num_async_tasks1=6,
+            expected_num_queries2=1286,
+            expected_num_async_tasks2=17,
+            expected_num_queries3=874,
+            expected_num_async_tasks3=16,
+        )
+
+    @override_settings(ENABLE_AUDITLOG=True)
+    def test_import_reimport_reimport_performance_pghistory_no_async(self):
+        """
+        This test checks the performance of the importers when using django-pghistory with async disabled.
+        Query counts will need to be determined by running the test initially.
+        """
+        configure_audit_system()
+        configure_pghistory_triggers()
+
+        testuser = User.objects.get(username="admin")
+        testuser.usercontactinfo.block_execution = True
+        testuser.usercontactinfo.save()
+
+        self._import_reimport_performance(
+            expected_num_queries1=1098,
+            expected_num_async_tasks1=6,
+            expected_num_queries2=1293,
+            expected_num_async_tasks2=17,
+            expected_num_queries3=881,
+            expected_num_async_tasks3=16,
+        )
+
+    @override_settings(ENABLE_AUDITLOG=True)
+    def test_import_reimport_reimport_performance_pghistory_no_async_with_product_grading(self):
+        """
+        This test checks the performance of the importers when using django-pghistory with async disabled and product grading enabled.
+        Query counts will need to be determined by running the test initially.
+        """
+        configure_audit_system()
+        configure_pghistory_triggers()
+
+        testuser = User.objects.get(username="admin")
+        testuser.usercontactinfo.block_execution = True
+        testuser.usercontactinfo.save()
+        self.system_settings(enable_product_grade=True)
+
+        self._import_reimport_performance(
+            expected_num_queries1=1105,
+            expected_num_async_tasks1=8,
+            expected_num_queries2=1300,
+            expected_num_async_tasks2=19,
+            expected_num_queries3=885,
+            expected_num_async_tasks3=18,
+        )
+
+    @override_settings(ENABLE_AUDITLOG=True)
+    def test_deduplication_performance_pghistory_async(self):
+        """Test deduplication performance with django-pghistory and async tasks enabled."""
+        configure_audit_system()
+        configure_pghistory_triggers()
+
+        self.system_settings(enable_deduplication=True)
+
+        self._deduplication_performance(
+            expected_num_queries1=264,
+            expected_num_async_tasks1=7,
+            expected_num_queries2=175,
+            expected_num_async_tasks2=7,
+            check_duplicates=False,  # Async mode - deduplication happens later
+        )
+
+    @override_settings(ENABLE_AUDITLOG=True)
+    def test_deduplication_performance_pghistory_no_async(self):
+        """Test deduplication performance with django-pghistory and async tasks disabled."""
+        configure_audit_system()
+        configure_pghistory_triggers()
+
+        self.system_settings(enable_deduplication=True)
+
+        testuser = User.objects.get(username="admin")
+        testuser.usercontactinfo.block_execution = True
+        testuser.usercontactinfo.save()
+
+        self._deduplication_performance(
+            expected_num_queries1=271,
+            expected_num_async_tasks1=7,
+            expected_num_queries2=236,
+            expected_num_async_tasks2=7,
+        )

From b5582d9942656dccdf7e5810c0ef9e9e10d0504c Mon Sep 17 00:00:00 2001
From: Valentijn Scholten <valentijnscholten@gmail.com>
Date: Fri, 27 Feb 2026 11:24:42 +0100
Subject: [PATCH 2/5] Add performance test query count update instructions to
 CONTRIBUTING.md

---
 readme-docs/CONTRIBUTING.md | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/readme-docs/CONTRIBUTING.md b/readme-docs/CONTRIBUTING.md
index 4c0b8a2210c..39992f7f772 100644
--- a/readme-docs/CONTRIBUTING.md
+++ b/readme-docs/CONTRIBUTING.md
@@ -55,6 +55,24 @@ Please use [these test scripts](../tests) to test your changes. These are the sc
 
 For changes that require additional settings, you can now use local_settings.py file. See the logging section below for more information.
 
+## Updating Performance Test Query Counts
+
+The importer performance tests in `unittests/test_importers_performance.py` assert on expected database query and async task counts. If your changes affect import behavior (e.g., adding queries or changing celery task usage), these counts may need to be updated.
+
+Run the update script to refresh expected counts:
+
+```bash
+python3 scripts/update_performance_test_counts.py
+```
+
+The script runs both `TestDojoImporterPerformanceSmall` (v2 endpoints) and `TestDojoImporterPerformanceSmallLocations` (v3 locations), captures actual counts, and updates the test file when they differ from expectations.
+
+To verify all tests pass after updating:
+
+```bash
+python3 scripts/update_performance_test_counts.py --verify
+```
+
 ## Python3 Version
 For compatibility reasons, the code in dev branch should be python3.13 compliant.
 

From 849258eaea2190e99dc56486da8dc1ec4c87905e Mon Sep 17 00:00:00 2001
From: Valentijn Scholten <valentijnscholten@gmail.com>
Date: Sat, 28 Feb 2026 23:48:53 +0100
Subject: [PATCH 3/5] Add _deduplication_performance to v3 Locations test and
 update expected counts

---
 unittests/test_importers_performance.py | 91 +++++++++++++++++++++++--
 1 file changed, 86 insertions(+), 5 deletions(-)

diff --git a/unittests/test_importers_performance.py b/unittests/test_importers_performance.py
index f8da76d1433..cc9c238335b 100644
--- a/unittests/test_importers_performance.py
+++ b/unittests/test_importers_performance.py
@@ -474,7 +474,7 @@ class TestDojoImporterPerformanceSmallLocations(TestDojoImporterPerformanceBase)
     Query counts are specific to the locations code path and will differ from the v2 endpoint counts.
 
     To determine or update the expected counts, run:
-        python3 scripts/update_performance_test_counts.py --test-class TestDojoImporterPerformanceSmallLocations
+        python3 scripts/update_performance_test_counts.py
     """
 
     def setUp(self):
@@ -571,6 +571,87 @@ def test_import_reimport_reimport_performance_pghistory_no_async_with_product_gr
             expected_num_async_tasks3=18,
         )
 
+    def _deduplication_performance(self, expected_num_queries1, expected_num_async_tasks1, expected_num_queries2, expected_num_async_tasks2, *, check_duplicates=True):
+        """
+        Test method to measure deduplication performance by importing the same scan twice.
+        Mirrors TestDojoImporterPerformanceSmall._deduplication_performance but uses
+        Locations-specific product/engagement names for test isolation.
+        """
+        _, engagement, lead, environment = self._create_test_objects(
+            "TestDojoDeduplicationPerformanceLocations",
+            "Test Deduplication Performance Engagement Locations",
+        )
+
+        with (  # noqa: SIM117
+            self.subTest("first_import"), impersonate(Dojo_User.objects.get(username="admin")),
+            STACK_HAWK_FILENAME.open(encoding="utf-8") as scan,
+        ):
+            with self.subTest(step="first_import", metric="queries"):
+                with self.assertNumQueries(expected_num_queries1):
+                    with self.subTest(step="first_import", metric="async_tasks"):
+                        with self._assertNumAsyncTask(expected_num_async_tasks1):
+                            import_options = {
+                                "user": lead,
+                                "lead": lead,
+                                "scan_date": None,
+                                "environment": environment,
+                                "minimum_severity": "Info",
+                                "active": True,
+                                "verified": True,
+                                "scan_type": STACK_HAWK_SCAN_TYPE,
+                                "engagement": engagement,
+                            }
+                            importer = DefaultImporter(**import_options)
+                            _, _, len_new_findings1, len_closed_findings1, _, _, _ = importer.process_scan(scan)
+
+        with (  # noqa: SIM117
+            self.subTest("second_import"), impersonate(Dojo_User.objects.get(username="admin")),
+            STACK_HAWK_FILENAME.open(encoding="utf-8") as scan,
+        ):
+            with self.subTest(step="second_import", metric="queries"):
+                with self.assertNumQueries(expected_num_queries2):
+                    with self.subTest(step="second_import", metric="async_tasks"):
+                        with self._assertNumAsyncTask(expected_num_async_tasks2):
+                            import_options = {
+                                "user": lead,
+                                "lead": lead,
+                                "scan_date": None,
+                                "environment": environment,
+                                "minimum_severity": "Info",
+                                "active": True,
+                                "verified": True,
+                                "scan_type": STACK_HAWK_SCAN_TYPE,
+                                "engagement": engagement,
+                            }
+                            importer = DefaultImporter(**import_options)
+                            _, _, len_new_findings2, len_closed_findings2, _, _, _ = importer.process_scan(scan)
+
+        logger.debug(f"First import: {len_new_findings1} new findings, {len_closed_findings1} closed findings")
+        logger.debug(f"Second import: {len_new_findings2} new findings, {len_closed_findings2} closed findings")
+
+        self.assertEqual(len_new_findings1, 6, "First import should create 6 new findings")
+        self.assertEqual(len_closed_findings1, 0, "First import should not close any findings")
+        self.assertEqual(len_new_findings2, 6, "Second import should report 6 new findings initially (before deduplication)")
+        self.assertEqual(len_closed_findings2, 0, "Second import should not close any findings")
+
+        if check_duplicates:
+            active_findings = Finding.objects.filter(
+                test__engagement=engagement,
+                active=True,
+                duplicate=False,
+            ).count()
+            duplicate_findings = Finding.objects.filter(
+                test__engagement=engagement,
+                duplicate=True,
+            ).count()
+            self.assertEqual(active_findings, 6, f"Expected 6 active findings, got {active_findings}")
+            self.assertEqual(duplicate_findings, 6, f"Expected 6 duplicate findings, got {duplicate_findings}")
+            total_findings = Finding.objects.filter(test__engagement=engagement).count()
+            self.assertEqual(total_findings, 12, f"Expected 12 total findings, got {total_findings}")
+        else:
+            total_findings = Finding.objects.filter(test__engagement=engagement).count()
+            self.assertEqual(total_findings, 12, f"Expected 12 total findings, got {total_findings}")
+
     @override_settings(ENABLE_AUDITLOG=True)
     def test_deduplication_performance_pghistory_async(self):
         """Test deduplication performance with django-pghistory and async tasks enabled."""
@@ -580,9 +661,9 @@ def test_deduplication_performance_pghistory_async(self):
         self.system_settings(enable_deduplication=True)
 
         self._deduplication_performance(
-            expected_num_queries1=264,
+            expected_num_queries1=1356,
             expected_num_async_tasks1=7,
-            expected_num_queries2=175,
+            expected_num_queries2=1165,
             expected_num_async_tasks2=7,
             check_duplicates=False,  # Async mode - deduplication happens later
         )
@@ -600,8 +681,8 @@ def test_deduplication_performance_pghistory_no_async(self):
         testuser.usercontactinfo.save()
 
         self._deduplication_performance(
-            expected_num_queries1=271,
+            expected_num_queries1=1363,
             expected_num_async_tasks1=7,
-            expected_num_queries2=236,
+            expected_num_queries2=1438,
             expected_num_async_tasks2=7,
         )

From 9ad75d0612d4a19716acdb62b2b10fd04a996f55 Mon Sep 17 00:00:00 2001
From: Valentijn Scholten <valentijnscholten@gmail.com>
Date: Sun, 1 Mar 2026 10:19:38 +0100
Subject: [PATCH 4/5] update

---
 scripts/update_performance_test_counts.py | 71 ++++++++++++-----------
 1 file changed, 38 insertions(+), 33 deletions(-)

diff --git a/scripts/update_performance_test_counts.py b/scripts/update_performance_test_counts.py
index 9b406539e6e..9b5a678ab63 100644
--- a/scripts/update_performance_test_counts.py
+++ b/scripts/update_performance_test_counts.py
@@ -14,15 +14,19 @@
     # Default: Update both v2 and v3 test classes
     python3 scripts/update_performance_test_counts.py
 
+    # Or update a specific test class:
+    python3 scripts/update_performance_test_counts.py --test-class TestDojoImporterPerformanceSmall
+    python3 scripts/update_performance_test_counts.py --test-class TestDojoImporterPerformanceSmallLocations
+
     # Step 1: Run tests and generate report only (without updating)
     python3 scripts/update_performance_test_counts.py --report-only
 
     # Step 2: Verify all tests pass
     python3 scripts/update_performance_test_counts.py --verify
 
-The script always runs and updates both TestDojoImporterPerformanceSmall (v2) and
-TestDojoImporterPerformanceSmallLocations (v3). The script defaults to --update
-behavior if no action flag is provided.
+By default (no --test-class) the script runs and updates both
+TestDojoImporterPerformanceSmall (v2) and TestDojoImporterPerformanceSmallLocations (v3).
+The script defaults to --update behavior if no action flag is provided.
 """
 
 import argparse
@@ -34,11 +38,11 @@
 # Path to the test file
 TEST_FILE = Path(__file__).parent.parent / "unittests" / "test_importers_performance.py"
 
-# Both v2 and v3 performance test classes - script always updates/verifies both
-TEST_CLASSES = (
+# All performance test classes, in run order
+TEST_CLASSES = [
     "TestDojoImporterPerformanceSmall",
     "TestDojoImporterPerformanceSmallLocations",
-)
+]
 
 
 class TestCount:
@@ -68,9 +72,9 @@ def extract_test_methods(test_class: str) -> list[str]:
 
     content = TEST_FILE.read_text()
 
-    # Find the test class definition (use (?<!-) to avoid matching "test-class" in docstrings)
+    # Find the test class definition
     class_pattern = re.compile(
-        rf"class {re.escape(test_class)}.*?(?=(?<!-)class \w|\Z)",
+        rf"class {re.escape(test_class)}.*?(?=class |\Z)",
         re.DOTALL,
     )
     class_match = class_pattern.search(content)
@@ -366,7 +370,7 @@ def generate_report(counts: list[TestCount], expected_counts: dict[str, dict[str
 
     print("=" * 80)
     print("\nTo update the test file, run:")
-    print("  python3 scripts/update_performance_test_counts.py")
+    print(f"  python scripts/update_performance_test_counts.py --test-class {test_name.split('_')[0]} --update")
     print()
 
 
@@ -374,8 +378,8 @@ def update_test_file(counts: list[TestCount], test_class: str | None = None):
     """
     Update the test file with new expected counts.
 
-    When test_class is provided, only update test methods within that class.
-    This is required when multiple classes (e.g. v2 and v3) share the same test method names.
+    When test_class is provided, method lookups are scoped to that class.
+    This is required because both v2 and v3 classes share the same method names.
     """
     if not counts:
         print("No counts to update.")
@@ -428,12 +432,13 @@ def _extract_call_span(method_content: str, call_name: str) -> tuple[int, int] |
         "second_import_async_tasks": "expected_num_async_tasks2",
     }
 
-    # Restrict search to the specified test class if given
+    # Restrict method search to the specified class to avoid updating the wrong
+    # class when v2 and v3 share identical method names.
     search_content = content
     search_offset = 0
     if test_class:
         class_pattern = re.compile(
-            rf"class {re.escape(test_class)}.*?(?=(?<!-)class \w|\Z)",
+            rf"class {re.escape(test_class)}.*?(?=class |\Z)",
             re.DOTALL,
         )
         class_match = class_pattern.search(content)
@@ -567,6 +572,12 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog=__doc__,
     )
+    parser.add_argument(
+        "--test-class",
+        required=False,
+        default=None,
+        help="Test class name to run (e.g., TestDojoImporterPerformanceSmall). Defaults to running all test classes if not provided.",
+    )
     parser.add_argument(
         "--report-only",
         action="store_true",
@@ -584,22 +595,21 @@ def main():
     )
 
     args = parser.parse_args()
+    classes_to_run = [args.test_class] if args.test_class else TEST_CLASSES
 
     if args.report_only:
-        # Step 1: Run tests and generate report for both test classes
         all_counts = []
-        for test_class in TEST_CLASSES:
+        for test_class in classes_to_run:
             test_methods = extract_test_methods(test_class)
             if not test_methods:
                 print(f"⚠️  No test methods found in {test_class}")
                 continue
 
-            print(f"\n{'=' * 80}")
-            print(f"Test class: {test_class} ({len(test_methods)} methods)")
+            print(f"\nFound {len(test_methods)} test method(s) in {test_class}")
             print("=" * 80)
 
             for test_method in test_methods:
-                print("\n---")
+                print(f"\n{'=' * 80}")
                 output, return_code = run_test_method(test_class, test_method)
                 success, error_msg = check_test_execution_success(output, return_code)
                 if not success:
@@ -611,35 +621,30 @@ def main():
                 if counts:
                     all_counts.extend(counts)
 
-        if all_counts:
-            generate_report(all_counts, {})
-        else:
-            print("✅ All tests passed! No count differences found.")
+            expected_counts = extract_expected_counts_from_file(test_class)
+            generate_report(all_counts, expected_counts)
 
     elif args.verify:
-        # Step 3: Verify both test classes
         all_pass = True
-        for test_class in TEST_CLASSES:
+        for test_class in classes_to_run:
             if not verify_tests(test_class):
                 all_pass = False
         sys.exit(0 if all_pass else 1)
 
     else:
         # Default: Update the file (--update is the default behavior)
-        # Run each test method in both test classes
         all_counts = []
-        for test_class in TEST_CLASSES:
+        for test_class in classes_to_run:
             test_methods = extract_test_methods(test_class)
             if not test_methods:
                 print(f"⚠️  No test methods found in {test_class}")
                 continue
 
-            print(f"\n{'=' * 80}")
-            print(f"Test class: {test_class} ({len(test_methods)} methods)")
+            print(f"\nFound {len(test_methods)} test method(s) in {test_class}")
             print("=" * 80)
 
             for test_method in test_methods:
-                print("\n---")
+                print(f"\n{'=' * 80}")
                 output, return_code = run_test_method(test_class, test_method)
                 success, error_msg = check_test_execution_success(output, return_code)
                 if not success:
@@ -654,7 +659,7 @@ def main():
 
                 if counts:
                     all_counts.extend(counts)
-                    # Update immediately after each test
+                    # Update immediately after each test, scoped to the current class
                     update_test_file(counts, test_class=test_class)
                     print(f"⚠️  {test_method}: Found {len(counts)} count mismatch(es) - updated file")
                 elif test_passed:
@@ -663,6 +668,7 @@ def main():
                     # Test might have failed for other reasons
                     print(f"⚠️  {test_method}: Test failed (exit code {return_code}) but no count mismatches parsed")
                     print("   This might indicate a parsing issue or a different type of failure")
+                    # Show a snippet of the output to help debug
                     fail_lines = [line for line in output.split("\n") if "FAIL" in line or "Error" in line or "Exception" in line]
                     if fail_lines:
                         print("   Relevant error lines:")
@@ -676,7 +682,7 @@ def main():
             # Do a final full-suite pass and apply any remaining mismatches so the suite passes as run in CI.
             print("\nRunning a final verify pass for stability...")
             all_pass = True
-            for test_class in TEST_CLASSES:
+            for test_class in classes_to_run:
                 success, suite_mismatches = verify_and_get_mismatches(test_class)
                 if not success and suite_mismatches:
                     print(f"\nApplying remaining mismatches from {test_class}...")
@@ -685,11 +691,10 @@ def main():
             if not all_pass:
                 print("\nRe-running verify...")
                 all_pass = True
-                for test_class in TEST_CLASSES:
+                for test_class in classes_to_run:
                     success, _ = verify_and_get_mismatches(test_class)
                     if not success:
                         all_pass = False
-                sys.exit(0 if all_pass else 1)
             sys.exit(0 if all_pass else 1)
         else:
             print(f"\n{'=' * 80}")

From a07f14aedd2ac483fe22c05097a20931e2f07843 Mon Sep 17 00:00:00 2001
From: Valentijn Scholten <valentijnscholten@gmail.com>
Date: Sun, 1 Mar 2026 10:44:20 +0100
Subject: [PATCH 5/5] cleanup

---
 run-unittest.sh | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/run-unittest.sh b/run-unittest.sh
index 9ffeeeb20b0..6073d4ff582 100755
--- a/run-unittest.sh
+++ b/run-unittest.sh
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash
 unset TEST_CASE
 unset FAIL_FAST
-unset KEEP_DB
-EXTRA_ARGS=()
 
 bash ./docker/docker-compose-check.sh
 if [[ $? -eq 1 ]]; then exit 1; fi
@@ -14,7 +12,6 @@ usage() {
   echo "Options:"
   echo "  --test-case -t {YOUR_FULLY_QUALIFIED_TEST_CASE}"
   echo "  --fail-fast -f - stop on first test failure"
-  echo "  --no-keepdb - recreate the test database (don't reuse existing)"
   echo "  --help -h - prints this dialogue."
   echo
   echo "You must specify a test case (arg)!"
@@ -44,10 +41,6 @@ while [[ $# -gt 0 ]]; do
       FAIL_FAST="--failfast"
       shift # past argument
       ;;
-    --no-keepdb)
-      KEEP_DB=""
-      shift # past argument
-      ;;
     -h|--help)
       usage
       exit 0
@@ -73,5 +66,4 @@ echo "Running docker compose unit tests with test case $TEST_CASE ..."
 # Compose V2 integrates compose functions into the Docker platform, continuing to support
 # most of the  previous docker-compose features and flags. You can run Compose V2 by
 # replacing the hyphen (-) with a space, using docker compose, instead of docker-compose.
-KEEP_DB="${KEEP_DB:---keepdb}"
-docker compose exec uwsgi bash -c "python manage.py test $TEST_CASE -v2 $KEEP_DB $FAIL_FAST ${EXTRA_ARGS[*]}"
+docker compose exec uwsgi bash -c "python manage.py test $TEST_CASE -v2 --keepdb $FAIL_FAST"