diff --git a/pgmq_sqlalchemy/queue.py b/pgmq_sqlalchemy/queue.py index 31fcaaa..5a59e44 100644 --- a/pgmq_sqlalchemy/queue.py +++ b/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 ( diff --git a/scripts/compelete_missing_test_for_queue.py b/scripts/compelete_missing_test_for_queue.py new file mode 100755 index 0000000..03db0ba --- /dev/null +++ b/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() diff --git a/scripts/scripts_utils/config.py b/scripts/scripts_utils/config.py index 61147cd..7259f0d 100644 --- a/scripts/scripts_utils/config.py +++ b/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" diff --git a/scripts/scripts_utils/operation_test_ast.py b/scripts/scripts_utils/operation_test_ast.py index f1b0e78..d0e83cd 100644 --- a/scripts/scripts_utils/operation_test_ast.py +++ b/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 diff --git a/scripts/scripts_utils/queue_ast.py b/scripts/scripts_utils/queue_ast.py index 7de5ac0..db8dc88 100644 --- a/scripts/scripts_utils/queue_ast.py +++ b/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("'''"): diff --git a/scripts/scripts_utils/test_ast.py b/scripts/scripts_utils/test_ast.py new file mode 100644 index 0000000..cfe3046 --- /dev/null +++ b/scripts/scripts_utils/test_ast.py @@ -0,0 +1,592 @@ +import libcst as cst +import re +import sys +from pathlib import Path +from typing import List, Set, Dict, Tuple + +sys.path.insert(0, str(Path(__file__).parent.parent.joinpath("scripts").resolve())) + + +class TestInfo: + """Information about a test function.""" + + def __init__(self, name: str, node: cst.FunctionDef): + self.name = name + self.node = node + self.is_test = name.startswith("test_") + self.is_async = name.endswith("_async") + self.base_name = name[:-6] if self.is_async else name + + +class ParseTestFunctionsVisitor(cst.CSTVisitor): + """CST visitor to parse test functions from test module""" + + def __init__(self): + self.tests: List[TestInfo] = [] + + def visit_FunctionDef(self, node: cst.FunctionDef): + func_name = node.name.value + if func_name.startswith("test_"): + self.tests.append(TestInfo(func_name, node)) + + +class AsyncTestTransformer(cst.CSTTransformer): + """Transform sync test to async test""" + + def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call: + """Transform method calls to async versions with await""" + if isinstance(updated_node.func, cst.Attribute): + # Check if calling method on pgmq object + if isinstance(updated_node.func.value, cst.Name): + var_name = updated_node.func.value.value + # Transform pgmq.method() to await pgmq.method_async() + if var_name == "pgmq": + method_name = updated_node.func.attr.value + # Add _async suffix to method name + new_attr = cst.Name(f"{method_name}_async") + new_func = updated_node.func.with_changes(attr=new_attr) + new_call = updated_node.with_changes(func=new_func) + # Wrap in await + return cst.Await(expression=new_call) + # Transform time.sleep() to await asyncio.sleep() + elif var_name == "time": + attr_name = updated_node.func.attr.value + if attr_name == "sleep": + # Change time.sleep to asyncio.sleep + new_func = updated_node.func.with_changes( + value=cst.Name("asyncio") + ) + new_call = updated_node.with_changes(func=new_func) + # Wrap in await + return cst.Await(expression=new_call) + + # Transform check_queue_exists to await check_queue_exists_async + if isinstance(updated_node.func, cst.Name): + if updated_node.func.value == "check_queue_exists": + 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 + + def leave_Param( + self, original_node: cst.Param, updated_node: cst.Param + ) -> cst.Param: + """Transform function parameters to use async fixtures.""" + param_name = updated_node.name.value + + 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 + + def leave_Name(self, original_node: cst.Name, updated_node: cst.Name) -> cst.Name: + """Transform variable name references to use async fixture names.""" + name = updated_node.value + + # Transform fixture references in function body + if name == "pgmq_setup_teardown": + return updated_node.with_changes(value="async_pgmq_setup_teardown") + elif name == "pgmq_partitioned_setup_teardown": + return updated_node.with_changes( + value="async_pgmq_partitioned_setup_teardown" + ) + elif name == "get_session_maker": + return updated_node.with_changes(value="get_async_session_maker") + elif name == "db_session": + return updated_node.with_changes(value="async_db_session") + + return updated_node + + def leave_FunctionDef( + self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef + ) -> cst.FunctionDef: + """Transform function to async and update name""" + # Transform function to async + new_node = updated_node.with_changes( + asynchronous=cst.Asynchronous(), + name=cst.Name(f"{updated_node.name.value}_async"), + ) + + # Transform docstring if exists + if updated_node.body.body and isinstance( + updated_node.body.body[0], cst.SimpleStatementLine + ): + first_stmt = updated_node.body.body[0] + if first_stmt.body and isinstance(first_stmt.body[0], cst.Expr): + expr = first_stmt.body[0] + if isinstance(expr.value, (cst.SimpleString, cst.ConcatenatedString)): + # Extract docstring value + if isinstance(expr.value, cst.SimpleString): + docstring = expr.value.value + else: + # For concatenated strings, skip transformation + docstring = None + + if docstring: + # Remove quotes to get actual string content + if docstring.startswith('"""') or docstring.startswith("'''"): + quote = docstring[:3] + content = docstring[3:-3] + elif docstring.startswith('"') or docstring.startswith("'"): + quote = docstring[0] + content = docstring[1:-1] + else: + content = docstring + quote = '"""' + + transformed_content = self.transform_docstring(content) + new_docstring = f"{quote}{transformed_content}{quote}" + + # Create new docstring node + new_expr = expr.with_changes( + value=cst.SimpleString(new_docstring) + ) + new_first_stmt = first_stmt.with_changes(body=[new_expr]) + + # Update body with new docstring + new_body = [new_first_stmt] + list(updated_node.body.body[1:]) + new_node = new_node.with_changes( + body=new_node.body.with_changes(body=new_body) + ) + + return new_node + + def transform_docstring(self, docstring: str) -> str: + """Transform docstring for async test version.""" + # Replace pgmq.method( with await pgmq.method_async( + modified = re.sub(r"(pgmq\.)(\w+)\(", r"await \1\2_async(", docstring) + # Replace time.sleep with await asyncio.sleep + modified = re.sub(r"time\.sleep\(", r"await asyncio.sleep(", modified) + return modified + + +class FillMissingTestsToModule(cst.CSTTransformer): + """CST Transformer to fill missing async tests to test module""" + + def __init__(self, to_add_async_tests: Dict[str, TestInfo]): + self.to_add_async_tests = to_add_async_tests + self.body_statements = [] + self.has_asyncio_import = False + self.has_check_queue_exists_async_import = False + + def visit_Module(self, node: cst.Module) -> bool: + """Collect all statements and check for asyncio import""" + self.body_statements = list(node.body) + # Check if asyncio is already imported + for stmt in node.body: + if isinstance(stmt, cst.SimpleStatementLine): + for item in stmt.body: + if isinstance(item, cst.Import): + for name in item.names: + if isinstance(name, cst.ImportAlias): + if ( + isinstance(name.name, cst.Name) + and name.name.value == "asyncio" + ): + self.has_asyncio_import = True + break + # Check for check_queue_exists_async import + if isinstance(item, cst.ImportFrom): + if ( + isinstance(item.module, cst.Attribute) + and isinstance(item.module.value, cst.Name) + and item.module.value.value == "tests" + and isinstance(item.module.attr, cst.Name) + and item.module.attr.value == "_utils" + ): + if isinstance(item.names, cst.ImportStar): + self.has_check_queue_exists_async_import = True + else: + for name in item.names: + if isinstance(name, cst.ImportAlias): + if ( + isinstance(name.name, cst.Name) + and name.name.value + == "check_queue_exists_async" + ): + self.has_check_queue_exists_async_import = ( + True + ) + break + return True + + def leave_Module( + self, original_node: cst.Module, updated_node: cst.Module + ) -> cst.Module: + """Add async tests after their sync counterparts and add asyncio import if needed""" + new_body = [] + inserted_asyncio = False + updated_check_queue_exists_import = False + updated_fixture_imports = False + updated_use_fixtures = False + seen_imports_from_utils = False + + for stmt in updated_node.body: + # Update fixture imports from tests.fixture_deps + if not updated_fixture_imports and isinstance( + stmt, cst.SimpleStatementLine + ): + for item in stmt.body: + if isinstance(item, cst.ImportFrom): + if ( + isinstance(item.module, cst.Attribute) + and isinstance(item.module.value, cst.Name) + and item.module.value.value == "tests" + and isinstance(item.module.attr, cst.Name) + and item.module.attr.value == "fixture_deps" + ): + # Add async fixture imports if needed + if not isinstance(item.names, cst.ImportStar): + has_pgmq_setup_teardown = False + has_async_pgmq_setup_teardown = False + has_pgmq_partitioned = False + has_async_pgmq_partitioned = False + + for name in item.names: + if isinstance(name, cst.ImportAlias) and isinstance( + name.name, cst.Name + ): + if name.name.value == "pgmq_setup_teardown": + has_pgmq_setup_teardown = True + elif ( + name.name.value + == "async_pgmq_setup_teardown" + ): + has_async_pgmq_setup_teardown = True + elif ( + name.name.value + == "pgmq_partitioned_setup_teardown" + ): + has_pgmq_partitioned = True + elif ( + name.name.value + == "async_pgmq_partitioned_setup_teardown" + ): + has_async_pgmq_partitioned = True + + # Add async versions if sync versions exist but async don't + new_names = list(item.names) + if ( + has_pgmq_setup_teardown + and not has_async_pgmq_setup_teardown + ): + new_names.append( + cst.ImportAlias( + name=cst.Name("async_pgmq_setup_teardown") + ) + ) + if ( + has_pgmq_partitioned + and not has_async_pgmq_partitioned + ): + new_names.append( + cst.ImportAlias( + name=cst.Name( + "async_pgmq_partitioned_setup_teardown" + ) + ) + ) + + if len(new_names) > len(item.names): + new_import = item.with_changes(names=new_names) + new_stmt = stmt.with_changes(body=[new_import]) + new_body.append(new_stmt) + updated_fixture_imports = True + continue + + # Update check_queue_exists import to include check_queue_exists_async + if isinstance(stmt, cst.SimpleStatementLine): + for item in stmt.body: + if isinstance(item, cst.ImportFrom): + if ( + isinstance(item.module, cst.Attribute) + and isinstance(item.module.value, cst.Name) + and item.module.value.value == "tests" + and isinstance(item.module.attr, cst.Name) + and item.module.attr.value == "_utils" + ): + # Skip duplicate imports from tests._utils + if seen_imports_from_utils: + continue + + # Check if check_queue_exists is imported + if not isinstance(item.names, cst.ImportStar): + has_check_queue_exists = False + has_check_queue_exists_async = False + for name in item.names: + if isinstance(name, cst.ImportAlias): + if isinstance(name.name, cst.Name): + if name.name.value == "check_queue_exists": + has_check_queue_exists = True + elif ( + name.name.value + == "check_queue_exists_async" + ): + has_check_queue_exists_async = True + + # If check_queue_exists is imported but not check_queue_exists_async, add it + if not updated_check_queue_exists_import: + if ( + has_check_queue_exists + and not has_check_queue_exists_async + ): + new_names = list(item.names) + new_names.append( + cst.ImportAlias( + name=cst.Name( + "check_queue_exists_async" + ) + ) + ) + new_import = item.with_changes(names=new_names) + new_stmt = stmt.with_changes(body=[new_import]) + new_body.append(new_stmt) + updated_check_queue_exists_import = True + seen_imports_from_utils = True + continue + else: + # Import already has both or is fine as-is + new_body.append(stmt) + seen_imports_from_utils = True + continue + + # Insert asyncio import after other imports if needed + if not inserted_asyncio and not self.has_asyncio_import: + if isinstance(stmt, cst.SimpleStatementLine): + # Check if this is the last import statement + is_import = any( + isinstance(item, (cst.Import, cst.ImportFrom)) + for item in stmt.body + ) + if is_import: + # Look ahead to see if next statement is not an import + current_idx = list(updated_node.body).index(stmt) + if current_idx + 1 < len(updated_node.body): + next_stmt = list(updated_node.body)[current_idx + 1] + if isinstance(next_stmt, cst.SimpleStatementLine): + next_is_import = any( + isinstance(item, (cst.Import, cst.ImportFrom)) + for item in next_stmt.body + ) + if not next_is_import: + # This is the last import, add asyncio after it + new_body.append(stmt) + new_body.append( + cst.SimpleStatementLine( + body=[ + cst.Import( + names=[ + cst.ImportAlias( + name=cst.Name("asyncio") + ) + ] + ) + ] + ) + ) + inserted_asyncio = True + continue + else: + # This is the last statement, add asyncio after it + new_body.append(stmt) + new_body.append( + cst.SimpleStatementLine( + body=[ + cst.Import( + names=[ + cst.ImportAlias( + name=cst.Name("asyncio") + ) + ] + ) + ] + ) + ) + inserted_asyncio = True + continue + + # Update use_fixtures list to include async versions + if not updated_use_fixtures and isinstance(stmt, cst.SimpleStatementLine): + for item in stmt.body: + if isinstance(item, cst.Assign): + for target in item.targets: + if ( + isinstance(target.target, cst.Name) + and target.target.value == "use_fixtures" + ): + # Check if value is a list + if isinstance(item.value, cst.List): + has_pgmq_setup_teardown = False + has_async_pgmq_setup_teardown = False + has_pgmq_partitioned = False + has_async_pgmq_partitioned = False + + for elem in item.value.elements: + if isinstance(elem.value, cst.Name): + if ( + elem.value.value + == "pgmq_setup_teardown" + ): + has_pgmq_setup_teardown = True + elif ( + elem.value.value + == "async_pgmq_setup_teardown" + ): + has_async_pgmq_setup_teardown = True + elif ( + elem.value.value + == "pgmq_partitioned_setup_teardown" + ): + has_pgmq_partitioned = True + elif ( + elem.value.value + == "async_pgmq_partitioned_setup_teardown" + ): + has_async_pgmq_partitioned = True + + # Add async versions to the list + new_elements = list(item.value.elements) + if ( + has_pgmq_setup_teardown + and not has_async_pgmq_setup_teardown + ): + new_elements.append( + cst.Element( + value=cst.Name( + "async_pgmq_setup_teardown" + ) + ) + ) + if ( + has_pgmq_partitioned + and not has_async_pgmq_partitioned + ): + new_elements.append( + cst.Element( + value=cst.Name( + "async_pgmq_partitioned_setup_teardown" + ) + ) + ) + + if len(new_elements) > len(item.value.elements): + new_list = item.value.with_changes( + elements=new_elements + ) + new_assign = item.with_changes(value=new_list) + new_stmt = stmt.with_changes(body=[new_assign]) + new_body.append(new_stmt) + updated_use_fixtures = True + continue + + new_body.append(stmt) + # If this is a sync test function, check if we need to add async version after it + if isinstance(stmt, cst.FunctionDef): + func_name = stmt.name.value + if func_name in self.to_add_async_tests: + # Add blank line before async test + new_body.append(cst.EmptyLine(whitespace=cst.SimpleWhitespace(""))) + new_body.append(cst.EmptyLine(whitespace=cst.SimpleWhitespace(""))) + # Add the async version right after the sync version + new_body.append(self.to_add_async_tests[func_name].node) + + return updated_node.with_changes(body=new_body) + + +def parse_test_functions_from_module( + module_tree: cst.Module, +) -> Tuple[List[TestInfo], Set[str]]: + """ + Parse test functions from test module CST Tree + + Args: + module_tree: cst.Module + + Returns: + Tuple of all_tests, missing_async_set + """ + + analyzer = ParseTestFunctionsVisitor() + module_tree.visit(analyzer) + + # Categorize test functions + async_tests_set = set() + missing_async_set = set() + + for test_info in analyzer.tests: + if test_info.is_async: + async_tests_set.add(test_info.base_name) + + # Find missing async tests + for test_info in analyzer.tests: + if not test_info.is_async and test_info.base_name not in async_tests_set: + missing_async_set.add(test_info.base_name) + + return analyzer.tests, missing_async_set + + +def transform_to_async_test( + transformer: AsyncTestTransformer, test_info: TestInfo +) -> TestInfo: + """Transform a sync test to async test. + + Args: + transformer: AsyncTestTransformer instance for CST transformations + test_info: TestInfo object containing sync test metadata + + Returns: + TestInfo object with transformed async test function + """ + orig_sync_func_node = test_info.node + async_node = orig_sync_func_node.visit(transformer) + + return TestInfo(f"{test_info.base_name}_async", async_node) + + +def get_async_tests_to_add( + sync_tests: List[TestInfo], missing_async: Set[str] +) -> Dict[str, TestInfo]: + """Generate async test functions for missing tests""" + transformer = AsyncTestTransformer() + async_tests: Dict[str, TestInfo] = {} + + for test_info in sync_tests: + if test_info.name in missing_async and not test_info.is_async: + async_tests[test_info.name] = transform_to_async_test( + transformer, test_info + ) + + return async_tests + + +def fill_missing_tests_to_module( + module_tree: cst.Module, to_add_async_tests: Dict[str, TestInfo] +) -> cst.Module: + """Fill missing async tests to module""" + transformer = FillMissingTestsToModule(to_add_async_tests=to_add_async_tests) + return module_tree.visit(transformer) diff --git a/tests/_utils.py b/tests/_utils.py index 237b489..f125552 100644 --- a/tests/_utils.py +++ b/tests/_utils.py @@ -1,8 +1,13 @@ +from typing import TYPE_CHECKING + from sqlalchemy import text -from sqlalchemy.orm import Session + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession -def check_queue_exists(db_session: Session, queue_name: str) -> bool: +def check_queue_exists(db_session: "Session", queue_name: str) -> bool: row = db_session.execute( text( "SELECT queue_name FROM pgmq.list_queues() WHERE queue_name = :queue_name ;" @@ -10,3 +15,17 @@ def check_queue_exists(db_session: Session, queue_name: str) -> bool: {"queue_name": queue_name}, ).first() return row is not None and row[0] == queue_name + + +async def check_queue_exists_async( + async_db_session: "AsyncSession", queue_name: str +) -> bool: + row = ( + await async_db_session.execute( + text( + "SELECT queue_name FROM pgmq.list_queues() WHERE queue_name = :queue_name ;" + ), + {"queue_name": queue_name}, + ) + ).first() + return row is not None and row[0] == queue_name diff --git a/tests/conftest.py b/tests/conftest.py index 9bf6393..a682f26 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -55,6 +55,8 @@ def pytest_generate_tests(metafunc: pytest.Metafunc): "pgmq_by_async_dsn", "pgmq_by_async_engine", "pgmq_by_async_session_maker", + "pgmq_by_async_dsn_and_async_engine", + "pgmq_by_async_dsn_and_async_session_maker", ] # Determine which fixtures to use @@ -207,6 +209,12 @@ def pgmq_by_dsn_and_engine(get_dsn, get_engine): return pgmq +@pytest.fixture(scope="function") +def pgmq_by_async_dsn_and_async_engine(get_async_dsn, get_async_engine): + pgmq = PGMQueue(dsn=get_async_dsn, engine=get_async_engine) + return pgmq + + @pytest.fixture(scope="function") def pgmq_by_dsn_and_session_maker(get_dsn, get_session_maker): pgmq = PGMQueue(dsn=get_dsn, session_maker=get_session_maker) @@ -214,5 +222,16 @@ def pgmq_by_dsn_and_session_maker(get_dsn, get_session_maker): @pytest.fixture(scope="function") -def db_session(get_session_maker) -> Session: +def pgmq_by_async_dsn_and_async_session_maker(get_async_dsn, get_async_session_maker): + pgmq = PGMQueue(dsn=get_async_dsn, session_maker=get_async_session_maker) + return pgmq + + +@pytest.fixture(scope="function") +def db_session(get_session_maker) -> "Session": return get_session_maker() + + +@pytest.fixture(scope="function") +async def async_db_session(get_async_session_maker) -> "AsyncSession": + return await get_async_session_maker() diff --git a/tests/fixture_deps.py b/tests/fixture_deps.py index d691f8c..5bd38d2 100644 --- a/tests/fixture_deps.py +++ b/tests/fixture_deps.py @@ -1,14 +1,14 @@ import uuid -from typing import Tuple +from typing import Tuple, Generator from inspect import iscoroutinefunction import pytest from pgmq_sqlalchemy import PGMQueue -from tests.constant import ASYNC_DRIVERS -from tests._utils import check_queue_exists +from tests.constant import SYNC_DRIVERS, ASYNC_DRIVERS +from tests._utils import check_queue_exists, check_queue_exists_async -PGMQ_WITH_QUEUE = Tuple[PGMQueue, str] +PGMQ_WITH_QUEUE = Generator[Tuple[PGMQueue, str], None, None] @pytest.fixture(scope="function") @@ -33,6 +33,10 @@ def test_something(pgmq_all_variants): pytest.skip( reason=f"Skip sync test: {request.function.__name__}, as driver: {driver_from_cli} is async" ) + if driver_from_cli and (driver_from_cli in SYNC_DRIVERS and is_async_test): + pytest.skip( + reason=f"Skip async test: {request.function.__name__}, as driver: {driver_from_cli} is sync" + ) return request.getfixturevalue(request.param) @@ -64,6 +68,36 @@ def test_something(pgmq_setup_teardown): assert check_queue_exists(db_session, queue_name) is False +@pytest.fixture(scope="function") +async def async_pgmq_setup_teardown( + pgmq_all_variants: PGMQueue, async_db_session +) -> PGMQ_WITH_QUEUE: + """ + Fixture that provides a PGMQueue instance with a unique temporary queue with setup and teardown. + + Args: + pgmq_all_variants (PGMQueue): The PGMQueue instance (parametrized across all variants). + async_db_session (sqlalchemy.ext.asyncio.AsyncSession): The SQLAlchemy session object. + + Yields: + tuple[PGMQueue,str]: A tuple containing the PGMQueue instance and the name of the temporary queue. + + Usage: + def test_something(pgmq_setup_teardown): + pgmq, queue_name = pgmq_setup_teardown + # test code here + + """ + pgmq = pgmq_all_variants + queue_name = f"test_queue_{uuid.uuid4().hex}" + assert await check_queue_exists_async(async_db_session, queue_name) is False + await pgmq.create_queue_async(queue_name) + assert check_queue_exists_async(async_db_session, queue_name) is True + yield pgmq, queue_name + await pgmq.drop_queue_async(queue_name) + assert await check_queue_exists_async(async_db_session, queue_name) is False + + @pytest.fixture(scope="function") def pgmq_partitioned_setup_teardown( pgmq_all_variants: PGMQueue, db_session @@ -92,3 +126,33 @@ def test_something(pgmq_partitioned_setup_teardown): yield pgmq, queue_name pgmq.drop_queue(queue_name, partitioned=True) assert check_queue_exists(db_session, queue_name) is False + + +@pytest.fixture(scope="function") +async def async_pgmq_partitioned_setup_teardown( + pgmq_all_variants: PGMQueue, async_db_session +) -> PGMQ_WITH_QUEUE: + """ + Fixture that provides a PGMQueue instance with a unique temporary partitioned queue with setup and teardown. + + Args: + pgmq_all_variants (PGMQueue): The PGMQueue instance (parametrized across all variants). + async_db_session (sqlalchemy.ext.asyncio.AsyncSession): The SQLAlchemy session object. + + Yields: + tuple[PGMQueue,str]: A tuple containing the PGMQueue instance and the name of the temporary queue. + + Usage: + def test_something(pgmq_partitioned_setup_teardown): + pgmq, queue_name = pgmq_partitioned_setup_teardown + # test code here + + """ + pgmq: PGMQueue = pgmq_all_variants + queue_name = f"test_queue_{uuid.uuid4().hex}" + assert await check_queue_exists_async(async_db_session, queue_name) is False + await pgmq.create_partitioned_queue(queue_name) + assert check_queue_exists_async(async_db_session, queue_name) is True + yield pgmq, queue_name + await pgmq.drop_queue(queue_name, partitioned=True) + assert await check_queue_exists_async(async_db_session, queue_name) is False