Add script to generate missing async tests using CST transformations

pgmq_sqlalchemy/queue.py

@@ -4,6 +4,7 @@
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.ext.asyncio import create_async_engine
 
+
 from .schema import Message, QueueMetrics
 from ._types import ENGINE_TYPE, SESSION_TYPE
 from ._utils import (

scripts/compelete_missing_test_for_queue.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# /// script
+# requires-python = ">=3.10,<3.11"
+# dependencies = [
+#   "rich>=13.6.0",
+#   "libcst>=1.0.0",
+#   "ruff>=0.14.10",
+# ]
+# ///
+"""
+Script to check for missing async test functions in test_queue.py and generate them.
+
+For each test function (starting with 'test_'), checks if there's a corresponding
+async test with the same name plus '_async' suffix. If missing, generates it.
+"""
+
+import libcst as cst
+import sys
+from pathlib import Path
+import contextlib
+import shutil
+
+import tempfile
+
+
+from scripts_utils.config import TEST_QUEUE_FILE, TEST_QUEUE_BACKUP_FILE
+from scripts_utils.console import console, user_input
+from scripts_utils.test_ast import (
+    parse_test_functions_from_module,
+    get_async_tests_to_add,
+    fill_missing_tests_to_module,
+)
+from scripts_utils.formatting import format_file, compare_file
+
+
+def main():
+    """Main function."""
+
+    module_tree = cst.parse_module(TEST_QUEUE_FILE.read_text())
+    sync_tests, missing_async = parse_test_functions_from_module(module_tree)
+
+    if not missing_async:
+        console.print(
+            "[bold green]SUCCESS:[/bold green] All test functions have corresponding async versions!"
+        )
+        sys.exit(0)
+
+    # log all the missing async tests
+    console.print()
+    console.print(
+        f"[bold yellow]WARNING:[/bold yellow] Found {len(missing_async)} missing async tests:",
+        style="bold",
+    )
+    for test in missing_async:
+        console.print(f"  [yellow]-[/yellow] {test}_async")
+    console.print()
+
+    # create missing async tests
+    async_tests_to_add = get_async_tests_to_add(sync_tests, missing_async)
+    # insert back to module
+    module_tree = fill_missing_tests_to_module(module_tree, async_tests_to_add)
+
+    # write back to tmp file for comparison
+    with tempfile.NamedTemporaryFile(mode="w+t", delete=False, suffix=".py") as f:
+        f.write(module_tree.code)
+        f.flush()
+        tmp_file = f.name
+        console.log(f"Complete missing async tests at {tmp_file}")
+
+    if tmp_file:
+        max_formatting = 3
+        for _ in range(max_formatting):
+            if format_file(tmp_file):
+                break
+
+    _, missing_async_for_tmp = parse_test_functions_from_module(
+        cst.parse_module(Path(tmp_file).read_text())
+    )
+
+    if missing_async_for_tmp:
+        console.log(
+            f"[error]Still have async tests missing after generating in {tmp_file}: {missing_async_for_tmp}[/]"
+        )
+    else:
+        console.log("[success]All missing async tests are generated[/]")
+
+    # compare existed test_queue.py and tmp.py
+    with contextlib.suppress(Exception):
+        compare_file(TEST_QUEUE_FILE, tmp_file)
+
+    # ask whether to apply the change
+    if user_input(f"Do you want to apply change to {TEST_QUEUE_FILE}"):
+        console.log(f"Backup existed {TEST_QUEUE_FILE} at {TEST_QUEUE_BACKUP_FILE}")
+        shutil.copy(TEST_QUEUE_FILE, TEST_QUEUE_BACKUP_FILE)
+        shutil.copy(tmp_file, TEST_QUEUE_FILE)
+        console.log("Add missing async tests successfully")
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

scripts/scripts_utils/config.py

@@ -6,3 +6,6 @@
 QUEUE_BACKUP_FILE = SOURCE_PATH / "queue_backup.py"
 OPERATION_FILE = SOURCE_PATH / "operation.py"
 OPERATION_BACKUP_FILE = SOURCE_PATH / "operation_backup.py"
+TESTS_PATH = PROJECT_ROOT / "tests"
+TEST_QUEUE_FILE = TESTS_PATH / "test_queue.py"
+TEST_QUEUE_BACKUP_FILE = TESTS_PATH / "test_queue_backup.py"

scripts/scripts_utils/operation_test_ast.py

@@ -68,9 +68,16 @@ def leave_Param(
         """Transform function parameters to use async fixtures."""
         param_name = updated_node.name.value
 
-        # Replace get_session_maker with get_async_session_maker
         if param_name == "get_session_maker":
             return updated_node.with_changes(name=cst.Name("get_async_session_maker"))
+        if param_name == "db_session":
+            return updated_node.with_changes(name=cst.Name("async_db_session"))
+        elif param_name == "pgmq_setup_teardown":
+            return updated_node.with_changes(name=cst.Name("async_pgmq_setup_teardown"))
+        elif param_name == "pgmq_partitioned_setup_teardown":
+            return updated_node.with_changes(
+                name=cst.Name("async_pgmq_partitioned_setup_teardown")
+            )
 
         return updated_node
 
@@ -100,13 +107,34 @@ def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Cal
                     )
                     return updated_node.with_changes(func=new_func)
 
-        # Check if this is a get_session_maker() call
+        # Check if this is a function call
         if isinstance(updated_node.func, cst.Name):
             if updated_node.func.value == "get_session_maker":
                 # Replace with get_async_session_maker
                 return updated_node.with_changes(
                     func=cst.Name("get_async_session_maker")
                 )
+            elif updated_node.func.value == "check_queue_exists":
+                # Replace with check_queue_exists_async and update db_session arg
+                new_call = updated_node.with_changes(
+                    func=cst.Name("check_queue_exists_async")
+                )
+                # Also need to update the db_session argument to async_db_session
+                if updated_node.args:
+                    new_args = []
+                    for arg in updated_node.args:
+                        if (
+                            isinstance(arg.value, cst.Name)
+                            and arg.value.value == "db_session"
+                        ):
+                            new_args.append(
+                                arg.with_changes(value=cst.Name("async_db_session"))
+                            )
+                        else:
+                            new_args.append(arg)
+                    new_call = new_call.with_changes(args=new_args)
+                # Wrap in await
+                return cst.Await(expression=new_call)
 
         return updated_node
 

scripts/scripts_utils/queue_ast.py

@@ -73,7 +73,6 @@ def leave_FunctionDef(
                     else:
                         # For concatenated strings, we'll skip transformation for now
                         docstring = None
-
                     if docstring:
                         # Remove quotes to get actual string content
                         if docstring.startswith('"""') or docstring.startswith("'''"):