Skip to content

Commit b8b40cf

Browse files
Copilotjason810496
Copilot
and
jason810496
authored
Add script to generate missing async tests for PGMQOperation using CST (#39)
* Add pre-commit hooks and scripts for async method checks in PGMQueue Distinguish sync and async operations in PGMQueue - Introduced a pre-commit hook to check for missing async methods in PGMQueue. - Added scripts to identify and generate missing async methods. - Created utility functions for AST manipulation and method transformation. - Established configuration for project paths and console output. * Refactor code transformation scripts from ast to libcst (#35) * Initial plan * Refactor AST-based code to use libcst for better code transformation Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com> * Fix: only wrap call expressions in await, not literals Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com> * Ask whether to apply change * Apply missing async methods * Correct files for check-sync-async-method-for-queue pre-commit * Add check for operation as well --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com> Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> * Initial plan * Add compelete_missing_test_for_operation.py script with CST-based approach Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com> * Fix code review feedback: remove redundant code and use Tuple from typing Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com> * Finalize test_operation * Fix _check_pg_partman_ext naming --------- Co-authored-by: LIU ZHE YOU <zhu424.dev@gmail.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: jason810496 <68415893+jason810496@users.noreply.github.com>
1 parent 9e76606 commit b8b40cf

5 files changed

Lines changed: 817 additions & 39 deletions

File tree

‎pgmq_sqlalchemy/queue.py‎

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
from sqlalchemy.orm import sessionmaker
55
from sqlalchemy.ext.asyncio import create_async_engine
66

7-
87
from .schema import Message, QueueMetrics
98
from ._types import ENGINE_TYPE, SESSION_TYPE
109
from ._utils import (
@@ -107,27 +106,15 @@ def __init__(
107106
bind=self.engine, class_=get_session_type(self.engine)
108107
)
109108

110-
def _check_pgmq_ext(self) -> None:
111-
"""Check if the pgmq extension exists."""
112-
self._execute_operation(PGMQOperation.check_pgmq_ext, session=None, commit=True)
113-
114-
async def _check_pgmq_ext_async(self) -> None:
115-
"""Check if the pgmq extension exists (async version)."""
116-
await self._execute_async_operation(
117-
PGMQOperation.check_pgmq_ext_async, session=None, commit=True
118-
)
109+
async def _check_pg_partman_ext_async(self) -> None:
110+
"""Check if the pg_partman extension exists."""
111+
async with self.session_maker() as session:
112+
await PGMQOperation.check_pg_partman_ext_async(session=session, commit=True)
119113

120114
def _check_pg_partman_ext(self) -> None:
121115
"""Check if the pg_partman extension exists."""
122-
self._execute_operation(
123-
PGMQOperation.check_pg_partman_ext, session=None, commit=True
124-
)
125-
126-
async def _check_pg_partman_ext_async(self) -> None:
127-
"""Check if the pg_partman extension exists (async version)."""
128-
await self._execute_async_operation(
129-
PGMQOperation.check_pg_partman_ext_async, session=None, commit=True
130-
)
116+
with self.session_maker() as session:
117+
PGMQOperation.check_pg_partman_ext(session=session, commit=True)
131118

132119
def _execute_operation(
133120
self,
@@ -338,7 +325,7 @@ async def create_partitioned_queue_async(
338325
339326
"""
340327
# check if the pg_partman extension exists before creating a partitioned queue at runtime
341-
await self._check_pg_partman_ext_async()
328+
self._check_pg_partman_ext()
342329

343330
return await self._execute_async_operation(
344331
PGMQOperation.create_partitioned_queue_async,
@@ -447,7 +434,7 @@ async def drop_queue_async(
447434
"""
448435
# check if the pg_partman extension exists before dropping a partitioned queue at runtime
449436
if partitioned:
450-
await self._check_pg_partman_ext_async()
437+
self._check_pg_partman_ext()
451438

452439
return await self._execute_async_operation(
453440
PGMQOperation.drop_queue_async,

‎scripts/compelete_missing_test_for_operation.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
#!/usr/bin/env python
2+
# /// script
3+
# requires-python = ">=3.10,<3.11"
4+
# dependencies = [
5+
# "rich>=13.6.0",
6+
# "libcst>=1.0.0",
7+
# ]
8+
# ///
9+
"""
10+
Script to check for missing async tests in test_operation.py and generate them.
11+
12+
For each public sync test (test_*_sync), checks if there's a corresponding
13+
async test with _async suffix. If missing, generates it using CST transformations.
14+
"""
15+
16+
import libcst as cst
17+
import sys
18+
from pathlib import Path
19+
import contextlib
20+
import shutil
21+
import tempfile
22+
23+
24+
from scripts_utils.console import console, user_input
25+
from scripts_utils.formatting import format_file, compare_file
26+
from scripts_utils.operation_test_ast import (
27+
parse_test_functions_from_module,
28+
get_async_tests_to_add,
29+
fill_missing_tests_to_module,
30+
)
31+
32+
33+
def main():
34+
"""Main function."""
35+
36+
# Define test file path
37+
PROJECT_ROOT = Path(__file__).parent.parent
38+
TEST_FILE = PROJECT_ROOT / "tests" / "test_operation.py"
39+
TEST_BACKUP_FILE = PROJECT_ROOT / "tests" / "test_operation_backup.py"
40+
41+
if not TEST_FILE.exists():
42+
console.print(f"[bold red]ERROR:[/bold red] Test file not found: {TEST_FILE}")
43+
sys.exit(1)
44+
45+
module_tree = cst.parse_module(TEST_FILE.read_text())
46+
all_tests, missing_async = parse_test_functions_from_module(module_tree)
47+
48+
if not missing_async:
49+
console.print(
50+
"[bold green]SUCCESS:[/bold green] All sync tests have corresponding async versions!"
51+
)
52+
sys.exit(0)
53+
54+
# Log all the missing async tests
55+
console.print()
56+
console.print(
57+
f"[bold yellow]WARNING:[/bold yellow] Found {len(missing_async)} missing async tests:",
58+
style="bold",
59+
)
60+
for test_name in sorted(missing_async):
61+
async_name = test_name.replace("_sync", "_async")
62+
console.print(f" [yellow]-[/yellow] {async_name}")
63+
console.print()
64+
65+
# Create missing async tests
66+
async_tests_to_add = get_async_tests_to_add(all_tests, missing_async)
67+
68+
# Insert back to module
69+
module_tree = fill_missing_tests_to_module(module_tree, async_tests_to_add)
70+
71+
# Write back to tmp file for comparison
72+
tmp_file = ""
73+
with tempfile.NamedTemporaryFile(mode="w+t", delete=False, suffix=".py") as f:
74+
f.write(module_tree.code)
75+
f.flush()
76+
tmp_file = f.name
77+
console.log(f"Generated missing async tests at {tmp_file}")
78+
79+
if tmp_file:
80+
max_formatting = 3
81+
for _ in range(max_formatting):
82+
if format_file(tmp_file):
83+
break
84+
85+
# Verify that all async tests are now present
86+
_, missing_async_for_tmp = parse_test_functions_from_module(
87+
cst.parse_module(Path(tmp_file).read_text())
88+
)
89+
90+
if missing_async_for_tmp:
91+
console.log(
92+
f"[error]Still have missing async tests after generation in {tmp_file}: {missing_async_for_tmp}[/]"
93+
)
94+
else:
95+
console.log("[success]All missing async tests are generated[/]")
96+
97+
# Compare existing test file and tmp file
98+
with contextlib.suppress(Exception):
99+
compare_file(TEST_FILE, tmp_file)
100+
101+
# Ask whether to apply the change
102+
if user_input(f"Do you want to apply change to {TEST_FILE}"):
103+
console.log(f"Backup existing {TEST_FILE} at {TEST_BACKUP_FILE}")
104+
shutil.copy(TEST_FILE, TEST_BACKUP_FILE)
105+
shutil.copy(tmp_file, TEST_FILE)
106+
console.log("Added missing async tests successfully")
107+
108+
sys.exit(0)
109+
110+
111+
if __name__ == "__main__":
112+
main()

0 commit comments

Comments
 (0)