Adopt new patterns up to some tests passing

Author: AlanCoding

dispatcher/brokers/base.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from typing import Optional
 
 
 class BaseBroker:
@@ -9,7 +10,7 @@ async def connect(self): ...
     async def aprocess_notify(self, connected_callback=None): ...
 
     @abstractmethod
-    async def apublish_message(self, channel, payload=None) -> None: ...
+    async def apublish_message(self, channel: Optional[str] = None, message: str = '') -> None: ...
 
     @abstractmethod
     async def aclose(self) -> None: ...

dispatcher/brokers/pg_notify.py

@@ -107,43 +107,24 @@ async def aprocess_notify(self, connected_callback=None):
                 async for notify in connection.notifies():
                     yield notify.channel, notify.payload
 
-    async def apublish_message(self, channel: Optional[str] = None, payload=None) -> None:
+    async def apublish_message(self, channel: Optional[str] = None, message: str = '') -> None:
         connection = await self.get_connection()
         channel = self.get_publish_channel(channel)
 
         async with connection.cursor() as cur:
-            if not payload:
+            if not message:
                 await cur.execute(f'NOTIFY {channel};')
             else:
-                await cur.execute(f"NOTIFY {channel}, '{payload}';")
+                await cur.execute(f"NOTIFY {channel}, '{message}';")
+
+        logger.debug(f'Sent pg_notify message of {len(message)} chars to {channel}')
 
     async def aclose(self) -> None:
         if self._connection:
             await self._connection.close()
             self._connection = None
 
 
-class ConnectionSaver:
-    def __init__(self):
-        self._connection = None
-
-
-connection_save = ConnectionSaver()
-
-
-def connection_saver(**config):
-    """
-    This mimics the behavior of Django for tests and demos
-    Philosophically, this is used by an application that uses an ORM,
-    or otherwise has its own connection management logic.
-    Dispatcher does not manage connections, so this a simulation of that.
-    """
-    if connection_save._connection is None:
-        config['autocommit'] = True
-        connection_save._connection = SyncBroker.create_connection(**config)
-    return connection_save._connection
-
-
 class SyncBroker(PGNotifyBase):
     def __init__(
         self,
@@ -177,11 +158,49 @@ def publish_message(self, channel: Optional[str] = None, message: str = '') -> N
         channel = self.get_publish_channel(channel)
 
         with connection.cursor() as cur:
-            cur.execute('SELECT pg_notify(%s, %s);', (channel, message))
+            if message:
+                cur.execute('SELECT pg_notify(%s, %s);', (channel, message))
+            else:
+                cur.execute(f'NOTIFY {channel};')
 
-        logger.debug(f'Sent pg_notify message to {channel}')
+        logger.debug(f'Sent pg_notify message of {len(message)} chars to {channel}')
 
     def close(self) -> None:
         if self._connection:
             self._connection.close()
             self._connection = None
+
+
+class ConnectionSaver:
+    def __init__(self) -> None:
+        self._connection: Optional[psycopg.Connection] = None
+        self._async_connection: Optional[psycopg.AsyncConnection] = None
+
+
+connection_save = ConnectionSaver()
+
+
+def connection_saver(**config) -> psycopg.Connection:
+    """
+    This mimics the behavior of Django for tests and demos
+    Philosophically, this is used by an application that uses an ORM,
+    or otherwise has its own connection management logic.
+    Dispatcher does not manage connections, so this a simulation of that.
+    """
+    if connection_save._connection is None:
+        config['autocommit'] = True
+        connection_save._connection = SyncBroker.create_connection(**config)
+    return connection_save._connection
+
+
+async def async_connection_saver(**config) -> psycopg.AsyncConnection:
+    """
+    This mimics the behavior of Django for tests and demos
+    Philosophically, this is used by an application that uses an ORM,
+    or otherwise has its own connection management logic.
+    Dispatcher does not manage connections, so this a simulation of that.
+    """
+    if connection_save._async_connection is None:
+        config['autocommit'] = True
+        connection_save._async_connection = await AsyncBroker.create_connection(**config)
+    return connection_save._async_connection

dispatcher/config.py

@@ -7,8 +7,8 @@
 
 class DispatcherSettings:
     def __init__(self, config: dict) -> None:
-        self.brokers: dict = config.get('brokers', [])
-        self.producers: dict = config.get('producers', [])
+        self.brokers: dict = config.get('brokers', {})
+        self.producers: dict = config.get('producers', {})
         self.service: dict = config.get('service', {'max_workers': 3})
         self.publish: dict = config.get('publish', {})
         # TODO: firmly planned sections of config for later

dispatcher/control.py

@@ -29,7 +29,7 @@ def __init__(self, queuename, send_data, expected_replies):
     def _create_events(self):
         return SimpleNamespace(exit_event=asyncio.Event())
 
-    async def process_message(self, payload, broker=None, channel=None):
+    async def process_message(self, payload, producer=None, channel=None):
         self.received_replies.append(payload)
         if self.expected_replies and (len(self.received_replies) >= self.expected_replies):
             self.events.exit_event.set()

dispatcher/factories.py

@@ -68,20 +68,24 @@ def get_sync_broker(broker_name, broker_config) -> BaseBroker:
     return broker_module.SyncBroker(**broker_config)
 
 
-def get_sync_publisher_from_settings(publish_broker: Optional[str] = None, settings: LazySettings = global_settings) -> BaseBroker:
+def _get_publisher_broker_name(publish_broker: Optional[str] = None, settings: LazySettings = global_settings) -> str:
     if publish_broker:
-        pass
+        return publish_broker
     elif len(settings.brokers) == 1:
-        publish_broker = list(settings.brokers.keys())[0]
+        return list(settings.brokers.keys())[0]
     elif 'default_broker' in settings.publish:
-        publish_broker = settings.publish['default_broker']
+        return settings.publish['default_broker']
     else:
         raise RuntimeError(f'Could not determine which broker to publish with between options {list(settings.brokers.keys())}')
 
-    return get_sync_broker(publish_broker, settings.brokers[publish_broker])
 
+def get_sync_publisher_from_settings(publish_broker: Optional[str] = None, settings: LazySettings = global_settings, **overrides) -> BaseBroker:
+    publish_broker = _get_publisher_broker_name(publish_broker=publish_broker, settings=settings)
 
-def get_async_publisher_from_settings(settings: LazySettings = global_settings, **overrides) -> BaseBroker:
+    return get_sync_broker(publish_broker, settings.brokers[publish_broker], **overrides)
+
+
+def get_async_publisher_from_settings(publish_broker: Optional[str] = None, settings: LazySettings = global_settings, **overrides) -> BaseBroker:
     """
     An asynchronous publisher is the ideal choice for submitting control-and-reply actions.
     This returns an asyncio broker of the default publisher type.
@@ -90,5 +94,5 @@ def get_async_publisher_from_settings(settings: LazySettings = global_settings,
     For control-and-reply, this will contain only the reply_to channel, to not receive
     unrelated traffic.
     """
-    publish_broker = settings.publish['default_broker']
+    publish_broker = _get_publisher_broker_name(publish_broker=publish_broker, settings=settings)
     return get_async_broker(publish_broker, settings.brokers[publish_broker], **overrides)

dispatcher/producers/brokered.py

@@ -36,10 +36,10 @@ async def produce_forever(self, dispatcher) -> None:
         self.dispatcher = dispatcher
         async for channel, payload in self.broker.aprocess_notify(connected_callback=self.connected_callback):
             self.produced_count += 1
-            await dispatcher.process_message(payload, broker=self, channel=channel)
+            await dispatcher.process_message(payload, producer=self, channel=channel)
 
-    async def notify(self, channel: str, payload: Optional[str] = None) -> None:
-        await self.broker.apublish_message(channel, payload=payload)
+    async def notify(self, channel: str, payload: str = '') -> None:
+        await self.broker.apublish_message(channel=channel, message=payload)
 
     async def shutdown(self) -> None:
         if self.production_task:

tests/conftest.py

@@ -13,6 +13,8 @@
 
 from dispatcher.brokers.pg_notify import SyncBroker, AsyncBroker
 from dispatcher.registry import DispatcherMethodRegistry
+from dispatcher.config import temporary_settings
+from dispatcher.factories import from_settings
 
 
 # List of channels to listen on
@@ -24,7 +26,10 @@
 BASIC_CONFIG = {
     "brokers": {
         "pg_notify": {
-            "channels": CHANNELS
+            "channels": CHANNELS,
+            "config": {'conninfo': CONNECTION_STRING},
+            "sync_connection_factory": "dispatcher.brokers.pg_notify.connection_saver",
+            # "async_connection_factory": "dispatcher.brokers.pg_notify.async_connection_saver",
         }
     },
     "pool": {
@@ -37,7 +42,7 @@
 async def aconnection_for_test():
     conn = None
     try:
-        conn = await AsyncBroker.create_connection({'conninfo': CONNECTION_STRING})
+        conn = await AsyncBroker.create_connection(conninfo=CONNECTION_STRING, autocommit=True)
 
         # Make sure database is running to avoid deadlocks which can come
         # from using the loop provided by pytest asyncio
@@ -58,27 +63,31 @@ def conn_config():
 
 @pytest.fixture
 def pg_dispatcher() -> DispatcherMain:
-    return DispatcherMain(BASIC_CONFIG)
+    # We can not reuse the connection between tests
+    config = BASIC_CONFIG.copy()
+    config['brokers']['pg_notify'].pop('async_connection_factory')
+    return DispatcherMain(config)
+
+
+@pytest.fixture
+def test_setup():
+    with temporary_settings(BASIC_CONFIG):
+        yield
 
 
 @pytest_asyncio.fixture(loop_scope="function", scope="function")
-async def apg_dispatcher(conn_config) -> AsyncIterator[DispatcherMain]:
-    # need to make a new connection because it can not be same as publisher
-    async with aconnection_for_test() as conn:
-        config = BASIC_CONFIG.copy()
-        config['producers']['BrokeredProducer']['connection'] = conn
-        # We have to fill in the config so that replies can still be sent in workers
-        # the workers may establish a new psycopg connection
-        config['producers']['BrokeredProducer']['config'] = conn_config
-        try:
-            dispatcher = DispatcherMain(config)
-
-            await dispatcher.connect_signals()
-            await dispatcher.start_working()
-            await dispatcher.wait_for_producers_ready()
-
-            yield dispatcher
-        finally:
+async def apg_dispatcher(test_setup) -> AsyncIterator[DispatcherMain]:
+    dispatcher = None
+    try:
+        dispatcher = from_settings()
+
+        await dispatcher.connect_signals()
+        await dispatcher.start_working()
+        await dispatcher.wait_for_producers_ready()
+
+        yield dispatcher
+    finally:
+        if dispatcher:
             await dispatcher.shutdown()
             await dispatcher.cancel_tasks()
 
@@ -87,18 +96,18 @@ async def apg_dispatcher(conn_config) -> AsyncIterator[DispatcherMain]:
 async def pg_message(psycopg_conn) -> Callable:
     async def _rf(message, channel='test_channel'):
         broker = AsyncBroker(connection=psycopg_conn)
-        await broker.apublish_message(channel, message)
+        await broker.apublish_message(channel=channel, message=message)
     return _rf
 
 
 @pytest_asyncio.fixture(loop_scope="function", scope="function")
-async def pg_control() -> AsyncIterator[Control]:
+async def pg_control(test_setup) -> AsyncIterator[Control]:
     """This has to use a different connection from dispatcher itself
 
     because psycopg will pool async connections, meaning that submission
     for the control task would be blocked by the listening query of the dispatcher itself"""
     async with aconnection_for_test() as conn:
-        yield Control('test_channel', async_connection=conn)
+        yield Control(queue='test_channel')
 
 
 @pytest_asyncio.fixture(loop_scope="function", scope="function")

tests/integration/publish/__init__.py

@@ -24,7 +24,7 @@ async def wait_to_receive(dispatcher, ct, timeout=5.0, interval=0.05):
 async def test_run_lambda_function(apg_dispatcher, pg_message):
     assert apg_dispatcher.pool.finished_count == 0
 
-    clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait())
+    clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait(), name='test_lambda_clear_wait')
     await pg_message('lambda: "This worked!"')
     await asyncio.wait_for(clearing_task, timeout=3)
 
@@ -36,7 +36,7 @@ async def test_run_decorated_function(apg_dispatcher, conn_config):
     assert apg_dispatcher.pool.finished_count == 0
 
     clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait())
-    test_methods.print_hello.apply_async(config=conn_config)
+    test_methods.print_hello.apply_async()
     await asyncio.wait_for(clearing_task, timeout=3)
 
     assert apg_dispatcher.pool.finished_count == 1
@@ -186,11 +186,11 @@ async def test_task_discard(apg_dispatcher, pg_message):
 
 
 @pytest.mark.asyncio
-async def test_task_discard_in_task_definition(apg_dispatcher, conn_config):
+async def test_task_discard_in_task_definition(apg_dispatcher):
     assert apg_dispatcher.pool.finished_count == 0
 
     for i in range(10):
-        test_methods.sleep_discard.apply_async(args=[2], config=conn_config)
+        test_methods.sleep_discard.apply_async(args=[2])
 
     await wait_to_receive(apg_dispatcher, 10)
 
@@ -199,11 +199,11 @@ async def test_task_discard_in_task_definition(apg_dispatcher, conn_config):
 
 
 @pytest.mark.asyncio
-async def test_tasks_in_serial(apg_dispatcher, conn_config):
+async def test_tasks_in_serial(apg_dispatcher):
     assert apg_dispatcher.pool.finished_count == 0
 
     for i in range(10):
-        test_methods.sleep_serial.apply_async(args=[2], config=conn_config)
+        test_methods.sleep_serial.apply_async(args=[2])
 
     await wait_to_receive(apg_dispatcher, 10)
 
@@ -212,11 +212,11 @@ async def test_tasks_in_serial(apg_dispatcher, conn_config):
 
 
 @pytest.mark.asyncio
-async def test_tasks_queue_one(apg_dispatcher, conn_config):
+async def test_tasks_queue_one(apg_dispatcher):
     assert apg_dispatcher.pool.finished_count == 0
 
     for i in range(10):
-        test_methods.sleep_queue_one.apply_async(args=[2], config=conn_config)
+        test_methods.sleep_queue_one.apply_async(args=[2])
 
     await wait_to_receive(apg_dispatcher, 10)