rewrite contextmanager arguments based on pytest output using redbaron

fix_num_db_queries_from_test_output.py

"""A script to help update many querycount assertions in code quickly.

Requirements:

    pip install redbaron

Deprecation note:

    redbaron (based on baron) is somewhat unmaintained and only supports
    Python grammar u to version 3.7, so for newer Python versions we may
    want to switch to something like parso:
    - https://github.com/PyCQA/baron#state-of-the-project
    - https://github.com/davidhalter/parso

Use like this:

- Run tests
- Some of them may produce output like (e.g. test summary):

    FAILED myproject/api/tests/test_video_algorithm.py::TestSerializeAlgorithm::test_large_threshold_single_page - Failed: Expected to perform 9 queries but 12 were done

- Validate that the updated querycounts are okay and we can use them.
- Copy-paste that test output to the TEST_OUTPUT string at the bottom of this file.
- Run this script like:

    python scripts/fix_num_db_queries_from_test_output.py

- The querycounts in tests code should have updated.
- Commit changes in git.

Note: Test summary may get truncated with longer lines, make sure to
    copy the output from individual tracebacks then so each line includes
    the actual querycount-done value.

"""

import os
import re

from redbaron import RedBaron

contextmanager_to_pattern = {
  "django_assert_num_queries": re.compile(
        r"FAILED (?P<testpath>.*) - Failed: Expected to perform (?P<num_expected>\d+) queries but (?P<num_done>\d+) were done"
    ),
}

testpath_pattern_with_cls = re.compile(r"(?P<rel_path>.*?)::(?P<cls>.*?)::(?P<func>.*)")
testpath_pattern_only_func = re.compile(r"(?P<rel_path>.*?)::(?P<func>.*)")

project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def main(test_output: str):
    for contextmanager_name, pattern in contextmanager_to_pattern.items():
        matches = pattern.findall(test_output)

        if contextmanager_name == "django_assert_num_queries":
            for testpath, num_expected, num_done in matches:

                match = testpath_pattern_with_cls.match(testpath)
                if match:
                    rel_path, cls, func = match.groups()
                else:
                    match = testpath_pattern_only_func.match(testpath)
                    if not match:
                        continue
                    cls = None
                    rel_path, func = match.groups()

                # Sanity check: We only want to affect test_*.py files for now.
                if not "test_" in rel_path or not rel_path.endswith(".py"):
                    raise ValueError("Not a test file: " + rel_path)

                filepath = os.path.join(project_root, rel_path)
                with open(filepath) as f:
                    code_before = f.read()

                red = RedBaron(code_before)
                test_function = None
                for node in red.filtered():
                    # Found the test function without any test class.
                    if node.name == func:
                        test_function = node
                        break
                    # Found the test class.
                    if node.name == cls:
                        if not hasattr(node, "filtered"):
                            continue
                        for childnode in node.filtered():
                            # Found the test function on the class.
                            if childnode.name == func:
                                test_function = childnode
                                break
                if not test_function:
                    continue

                expected_contextmanager_code = (
                    f"with {contextmanager_name}({num_expected}):"
                )
                for node in test_function.filtered():
                    if node.__class__.__name__ == "WithNode" and str(node).startswith(
                        expected_contextmanager_code
                    ):
                        # Hint: Debug via <node>.help(deep=True)
                        for contextitemnode in node.contexts:
                            for callargumentnode in contextitemnode.value.call:
                                intnode = callargumentnode.value
                                if intnode.value == num_expected:
                                    intnode.value = num_done
                                    print(
                                        "FIXED:",
                                        cls,
                                        func,
                                        contextmanager_name,
                                        num_expected,
                                        "->",
                                        num_done,
                                    )
                                    break

                # Update the file.
                code_after = red.dumps()
                with open(filepath, "w") as f:
                    f.write(code_after)


TEST_OUTPUT = """

FAILED myproject/api/tests/test_video_algorithm.py::TestSerializeAlgorithm::test_tiny_threshold_many_pages - Failed: Expected to perform 16 queries but 19 were done
FAILED myproject/api/tests/test_video.py::TestGenerateVideos::test_changing_overflow_threshold_creates_video - Failed: Expected to perform 19 queries but 22 were done

"""


if __name__ == "__main__":
    main(TEST_OUTPUT)