Route outgoing messages to correct channel (#99)

If it is part of a transaction then use the relevant transacted channel,
otherwise use the (untransacted) shared channel.

Resolves #98

Author: Anthchirp

src/workflows/transport/pika_transport.py

@@ -438,7 +438,14 @@ def _unsubscribe(self, sub_id: int, **kwargs):
         # Callback reference is kept as further messages may already have been received
 
     def _send(
-        self, destination, message, headers=None, delay=None, expiration=None, **kwargs
+        self,
+        destination,
+        message,
+        headers=None,
+        delay=None,
+        expiration=None,
+        transaction: Optional[int] = None,
+        **kwargs,
     ):
         """
         Send a message to a queue.
@@ -449,6 +456,7 @@ def _send(
             headers: Further arbitrary headers to pass to pika
             delay: Delay transport of message by this many seconds
             expiration: Optional TTL expiration time, relative to sending time
+            transaction: Transaction ID if message should be part of a transaction
         """
         if not headers:
             headers = {}
@@ -470,6 +478,7 @@ def _send(
             body=message,
             properties=properties,
             mandatory=True,
+            transaction_id=transaction,
         ).result()
 
     def _broadcast(
@@ -479,6 +488,7 @@ def _broadcast(
         headers=None,
         delay=None,
         expiration: Optional[int] = None,
+        transaction: Optional[int] = None,
         **kwargs,
     ):
         """Send a message to a fanout exchange.
@@ -489,6 +499,7 @@ def _broadcast(
             headers: Further arbitrary headers to pass to pika
             delay: Delay transport of message by this many seconds
             expiration: Optional TTL expiration time, in seconds, relative to sending time
+            transaction: Transaction ID if message should be part of a transaction
             kwargs: Arbitrary arguments for other transports. Ignored.
         """
         assert delay is None, "Delay Not implemented"
@@ -511,6 +522,7 @@ def _broadcast(
             body=message,
             properties=properties,
             mandatory=False,
+            transaction_id=transaction,
         ).result()
 
     def _transaction_begin(
@@ -930,6 +942,7 @@ def send(
         body: Union[str, bytes],
         properties: pika.spec.BasicProperties = None,
         mandatory: bool = True,
+        transaction_id: Optional[int] = None,
     ) -> Future[None]:
         """Send a message. Thread-safe."""
 
@@ -941,7 +954,11 @@ def send(
         def _send():
             if future.set_running_or_notify_cancel():
                 try:
-                    self._get_shared_channel().basic_publish(
+                    if transaction_id:
+                        channel = self._transactions_by_id[transaction_id]
+                    else:
+                        channel = self._get_shared_channel()
+                    channel.basic_publish(
                         exchange=exchange,
                         routing_key=routing_key,
                         body=body,

tests/transport/test_pika.py

@@ -93,16 +93,16 @@ def test_check_config_file_behaviour(mockpika, mock_pikathread, tmp_path):
         cfgfile = tmp_path / "config"
         cfgfile.write_text(
             """
-    # An example pika configuration file
-    # Only lines in the [pika] block will be interpreted
-
-    [rabbit]
-    host = localhost
-    port = 5672
-    username = someuser
-    password = somesecret
-    vhost = namespace
-    """
+            # An example pika configuration file
+            # Only lines in the [pika] block will be interpreted
+
+            [rabbit]
+            host = localhost
+            port = 5672
+            username = someuser
+            password = somesecret
+            vhost = namespace
+            """
         )
 
         parser.parse_args(
@@ -218,6 +218,7 @@ def test_broadcast_status(mockpika, mock_pikathread):
         "body": mock.ANY,
         "properties": mock.ANY,
         "mandatory": False,
+        "transaction_id": None,
     }
     statusdict = json.loads(kwargs.get("body"))
     assert statusdict["status"] == str(mock.sentinel.status)
@@ -242,6 +243,7 @@ def test_send_message(mockpika, mock_pikathread):
         "body": mock.sentinel.message,
         "mandatory": True,
         "properties": mock.ANY,
+        "transaction_id": None,
     }
     assert mockproperties.call_args[1].get("headers") == {}
     assert int(mockproperties.call_args[1].get("delivery_mode")) == 2
@@ -291,6 +293,7 @@ def test_sending_message_with_expiration(mockpika, mock_pikathread):
         "body": mock.sentinel.message,
         "mandatory": True,
         "properties": mock.ANY,
+        "transaction_id": None,
     }
     assert int(mockproperties.return_value.expiration) == 120 * 1000
 
@@ -341,6 +344,7 @@ def test_send_broadcast(mockpika, mock_pikathread):
         "body": mock.sentinel.message,
         "properties": mock.ANY,
         "mandatory": False,
+        "transaction_id": None,
     }
 
     transport._broadcast(
@@ -360,6 +364,7 @@ def test_send_broadcast(mockpika, mock_pikathread):
         "body": mock.sentinel.message,
         "properties": mock.ANY,
         "mandatory": False,
+        "transaction_id": None,
     }
 
     # Delay not implemented yet
@@ -404,6 +409,7 @@ def test_broadcasting_message_with_expiration(mockpika, mock_pikathread):
         "body": mock.sentinel.message,
         "properties": mock.ANY,
         "mandatory": False,
+        "transaction_id": None,
     }
 
 
@@ -447,6 +453,7 @@ def test_messages_are_serialized_for_transport(mock_pikathread):
         "body": banana_str,
         "properties": pika.BasicProperties(delivery_mode=2, headers={}),
         "mandatory": True,
+        "transaction_id": None,
     }
 
     transport.broadcast(str(mock.sentinel.queue2), banana)
@@ -458,6 +465,7 @@ def test_messages_are_serialized_for_transport(mock_pikathread):
         "body": banana_str,
         "properties": pika.BasicProperties(delivery_mode=2, headers={}),
         "mandatory": False,
+        "transaction_id": None,
     }
 
     with pytest.raises(TypeError):
@@ -481,6 +489,7 @@ def test_messages_are_not_serialized_for_raw_transport(_mockpika, mock_pikathrea
         "body": banana,
         "mandatory": True,
         "properties": mock.ANY,
+        "transaction_id": None,
     }
 
     mock_pikathread.send.reset_mock()
@@ -494,6 +503,7 @@ def test_messages_are_not_serialized_for_raw_transport(_mockpika, mock_pikathrea
         "body": banana,
         "properties": mock.ANY,
         "mandatory": False,
+        "transaction_id": None,
     }
 
     mock_pikathread.send.reset_mock()
@@ -507,6 +517,7 @@ def test_messages_are_not_serialized_for_raw_transport(_mockpika, mock_pikathrea
         "body": mock.sentinel.unserializable,
         "mandatory": True,
         "properties": mock.ANY,
+        "transaction_id": None,
     }
 
 
@@ -818,7 +829,6 @@ def temporary_queue_declare(
             ):
                 """
                 Declare an auto-named queue that is automatically deleted on test end.
-
                 """
                 queue = self.channel.queue_declare(
                     "", auto_delete=auto_delete, exclusive=exclusive, **kwargs
@@ -994,6 +1004,7 @@ def test_pikathread_send(connection_params, test_channel):
         ).result()
 
         # But should fail with it declared
+        pytest.xfail("UnroutableError is not raised without publisher confirms, #96")
         with pytest.raises(pika.exceptions.UnroutableError):
             thread.send(
                 "", "unroutable-missing-queue", "Another Message", mandatory=True