Fix wrong event loop usage from PGMQueue

Author: jason810496

pgmq_sqlalchemy/queue.py

@@ -1,4 +1,3 @@
-import asyncio
 from typing import List, Optional, TYPE_CHECKING
 
 from sqlalchemy import create_engine
@@ -28,14 +27,12 @@ class PGMQueue:
 
     is_async: bool = False
     is_pg_partman_ext_checked: bool = False
-    loop: asyncio.AbstractEventLoop = None
 
     def __init__(
         self,
         dsn: Optional[str] = None,
         engine: Optional[ENGINE_TYPE] = None,
         session_maker: Optional[sessionmaker] = None,
-        loop: Optional[asyncio.AbstractEventLoop] = None,
     ) -> None:
         """
 
@@ -84,8 +81,6 @@ def __init__(
             dsn (Optional[str]): Database connection string.
             engine (Optional[ENGINE_TYPE]): SQLAlchemy engine (sync or async).
             session_maker (Optional[sessionmaker]): SQLAlchemy session maker.
-            loop (Optional[asyncio.AbstractEventLoop]): Event loop for async operations.
-                If not provided, a new event loop will be created for async engines.
 
         .. note::
             | ``PGMQueue`` will **auto create** the ``pgmq`` extension ( and ``pg_partman`` extension if the method is related with **partitioned_queue** ) if it does not exist in the Postgres.
@@ -112,17 +107,6 @@ def __init__(
                 bind=self.engine, class_=get_session_type(self.engine)
             )
 
-        if self.is_async:
-            if loop is not None:
-                # Use the provided event loop
-                self.loop = loop
-            else:
-                # Create a new event loop
-                self.loop = asyncio.new_event_loop()
-
-        # create pgmq extension if not exists
-        self._check_pgmq_ext()
-
     def _check_pgmq_ext(self) -> None:
         """Check if the pgmq extension exists."""
         self._execute_operation(PGMQOperation.check_pgmq_ext, session=None, commit=True)

tests/conftest.py

@@ -10,10 +10,14 @@
 from tests.constant import ASYNC_DRIVERS, SYNC_DRIVERS
 
 # Async fixture names for test filtering
-ASYNC_FIXTURE_NAMES = ['pgmq_by_async_dsn', 'pgmq_by_async_engine', 'pgmq_by_async_session_maker']
+ASYNC_FIXTURE_NAMES = [
+    "pgmq_by_async_dsn",
+    "pgmq_by_async_engine",
+    "pgmq_by_async_session_maker",
+]
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser: pytest.Parser):
     """Add custom command-line options for pytest."""
     parser.addoption(
         "--driver",
@@ -29,30 +33,30 @@ def pytest_addoption(parser):
     )
 
 
-def pytest_generate_tests(metafunc):
+def pytest_generate_tests(metafunc: pytest.Metafunc):
     """
     Dynamically generate test parametrization based on CLI options.
-    
+
     This allows us to parametrize fixtures based on the --driver option.
     """
     if "pgmq_all_variants" in metafunc.fixturenames:
         driver_from_cli = metafunc.config.getoption("--driver")
-        
+
         # Define sync and async fixture variants
         sync_fixtures = [
-            'pgmq_by_dsn',
-            'pgmq_by_engine',
-            'pgmq_by_session_maker',
-            'pgmq_by_dsn_and_engine',
-            'pgmq_by_dsn_and_session_maker',
+            "pgmq_by_dsn",
+            "pgmq_by_engine",
+            "pgmq_by_session_maker",
+            "pgmq_by_dsn_and_engine",
+            "pgmq_by_dsn_and_session_maker",
         ]
-        
+
         async_fixtures = [
-            'pgmq_by_async_dsn',
-            'pgmq_by_async_engine',
-            'pgmq_by_async_session_maker',
+            "pgmq_by_async_dsn",
+            "pgmq_by_async_engine",
+            "pgmq_by_async_session_maker",
         ]
-        
+
         # Determine which fixtures to use
         if not driver_from_cli:
             # No driver specified, use all fixtures
@@ -63,13 +67,9 @@ def pytest_generate_tests(metafunc):
         else:
             # Sync driver specified
             fixture_params = sync_fixtures
-        
+
         # Parametrize the test
-        metafunc.parametrize(
-            "pgmq_all_variants",
-            fixture_params,
-            indirect=True
-        )
+        metafunc.parametrize("pgmq_all_variants", fixture_params, indirect=True)
 
 
 @pytest.fixture(scope="module")
@@ -93,7 +93,7 @@ def get_sa_password():
 
 
 @pytest.fixture(scope="module")
-def get_sa_db(request):
+def get_sa_db(request: pytest.FixtureRequest):
     """Get database name from CLI argument or environment variable."""
     db_name_from_cli = request.config.getoption("--db-name")
     if db_name_from_cli:
@@ -112,14 +112,14 @@ def get_dsn(
 ):
     """Get DSN for sync drivers based on CLI option."""
     driver_from_cli = request.config.getoption("--driver")
-    
+
     # Use CLI driver if specified and it's a sync driver
     if driver_from_cli and driver_from_cli in SYNC_DRIVERS:
         driver = driver_from_cli
     else:
         # Default to first sync driver if no CLI option or invalid
         driver = SYNC_DRIVERS[0]
-    
+
     return f"postgresql+{driver}://{get_sa_user}:{get_sa_password}@{get_sa_host}:{get_sa_port}/{get_sa_db}"
 
 
@@ -134,14 +134,14 @@ def get_async_dsn(
 ):
     """Get DSN for async drivers based on CLI option."""
     driver_from_cli = request.config.getoption("--driver")
-    
+
     # Use CLI driver if specified and it's an async driver
     if driver_from_cli and driver_from_cli in ASYNC_DRIVERS:
         driver = driver_from_cli
     else:
         # Default to first async driver if no CLI option or invalid
         driver = ASYNC_DRIVERS[0]
-    
+
     return f"postgresql+{driver}://{get_sa_user}:{get_sa_password}@{get_sa_host}:{get_sa_port}/{get_sa_db}"
 
 

tests/fixture_deps.py

@@ -1,9 +1,11 @@
 import uuid
 from typing import Tuple
+from inspect import iscoroutinefunction
 
 import pytest
 
 from pgmq_sqlalchemy import PGMQueue
+from tests.constant import ASYNC_DRIVERS
 from tests._utils import check_queue_exists
 
 PGMQ_WITH_QUEUE = Tuple[PGMQueue, str]
@@ -13,18 +15,24 @@
 def pgmq_all_variants(request: pytest.FixtureRequest) -> PGMQueue:
     """
     Fixture that parametrizes tests across all appropriate PGMQueue initialization methods.
-    
+
     When --driver is specified, only fixtures matching that driver type (sync/async) are used.
     Without --driver, all fixtures are used.
-    
+
     The parametrization is handled by pytest_generate_tests in conftest.py.
-    
+
     Usage:
         def test_something(pgmq_all_variants):
             pgmq: PGMQueue = pgmq_all_variants
             # test code here
     """
     # The param is set by pytest_generate_tests via indirect parametrization
+    is_async_test = iscoroutinefunction(request.function)
+    driver_from_cli = request.config.getoption("--driver")
+    if driver_from_cli and (driver_from_cli in ASYNC_DRIVERS and not is_async_test):
+        pytest.skip(
+            reason=f"Skip sync test: {request.function.__name__}, as driver: {driver_from_cli} is async"
+        )
     return request.getfixturevalue(request.param)
 
 

tests/test_construct_pgmq.py

@@ -3,6 +3,8 @@
 
 from tests.fixture_deps import pgmq_all_variants
 
+use_fixtures = [pgmq_all_variants]
+
 
 def test_construct_pgmq(pgmq_all_variants):
     pgmq: PGMQueue = pgmq_all_variants

tests/test_event_loop.py

@@ -8,18 +8,15 @@
 
 from tests.fixture_deps import (
     PGMQ_WITH_QUEUE,
-    pgmq_all_variants,
     pgmq_setup_teardown,
     pgmq_partitioned_setup_teardown,
+    pgmq_all_variants,
 )
 
 from tests._utils import check_queue_exists
 from tests.constant import MSG, LOCK_FILE_NAME
 
-use_fixtures = [
-    pgmq_setup_teardown,
-    pgmq_partitioned_setup_teardown,
-]
+use_fixtures = [pgmq_setup_teardown, pgmq_partitioned_setup_teardown, pgmq_all_variants]
 
 
 def test_create_queue(pgmq_all_variants, db_session):
@@ -486,13 +483,13 @@ def test_read_with_poll_without_vt(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
     """Test read_with_poll when vt parameter is not provided (None)."""
 
     pgmq, queue_name = pgmq_setup_teardown
-    
+
     # Set a custom default vt for the pgmq instance
     pgmq.vt = 100
-    
+
     # Send a message
     msg_id = pgmq.send(queue_name, MSG)
-    
+
     # Call read_with_poll with vt=None to test the fallback logic
     # When vt is None, it should fall back to using pgmq.vt value (100)
     msgs = pgmq.read_with_poll(
@@ -502,67 +499,43 @@ def test_read_with_poll_without_vt(pgmq_setup_teardown: PGMQ_WITH_QUEUE):
         max_poll_seconds=2,
         poll_interval_ms=100,
     )
-    
+
     assert msgs is not None
     assert len(msgs) == 1
     assert msgs[0].msg_id == msg_id
     assert msgs[0].message == MSG
 
 
-def test_execute_operation_with_provided_sync_session(pgmq_by_session_maker, get_session_maker, db_session):
+def test_execute_operation_with_provided_sync_session(
+    pgmq_by_session_maker, get_session_maker, db_session
+):
     """Test _execute_operation sync path when session is provided."""
 
     pgmq: PGMQueue = pgmq_by_session_maker
     queue_name = f"test_queue_{uuid.uuid4().hex}"
-    
+
     # Create a session to pass to the operations
     # Using the same session across multiple operations demonstrates
     # that the sync path with provided session works correctly
     with get_session_maker() as session:
         # Create queue with provided session
         pgmq.create_queue(queue_name, session=session)
-        
+
         # Verify queue was created
         assert check_queue_exists(db_session, queue_name) is True
-        
+
         # Send a message with the same provided session
         msg_id = pgmq.send(queue_name, MSG, session=session)
-        
+
         # Read message with the same provided session
         msg = pgmq.read(queue_name, vt=30, session=session)
-        
+
         assert msg is not None
         assert msg.msg_id == msg_id
         assert msg.message == MSG
-        
+
         # Clean up with the same provided session
         pgmq.drop_queue(queue_name, session=session)
-    
-    # Verify queue was dropped
-    assert check_queue_exists(db_session, queue_name) is False
-
 
-def test_execute_operation_async_with_session_none(pgmq_by_async_dsn, db_session):
-    """Test _execute_operation async path when session is None."""
-
-    pgmq: PGMQueue = pgmq_by_async_dsn
-    queue_name = f"test_queue_{uuid.uuid4().hex}"
-    
-    # Verify this is an async PGMQueue
-    assert pgmq.is_async is True
-    
-    # When session is None, _execute_operation creates a new async session
-    # and uses loop.run_until_complete to execute the operation
-    pgmq.create_queue(queue_name)
-    msg_id = pgmq.send(queue_name, MSG)
-    msg = pgmq.read(queue_name, vt=30)
-    
-    assert msg is not None
-    assert msg.msg_id == msg_id
-    assert msg.message == MSG
-    
-    # Clean up
-    pgmq.drop_queue(queue_name)
-    
     # Verify queue was dropped
     assert check_queue_exists(db_session, queue_name) is False