Update broadway.ex

lib/broadway.ex

@@ -1268,6 +1268,18 @@ defmodule Broadway do
       when set to `:flush`, the batch the message is in is immediately delivered. When set
       to `:bulk`, batch is delivered when its size or timeout is reached.
       """
+    ],
+    batcher: [
+      type: :atom,
+      default: :default,
+      doc: """
+      the name of the batcher assigned to the test message. This option is useful
+      when the producer's `:transform` option sets the batcher, as test messages
+      do not run through the producer. If `handle_message/3` sets the batcher, its
+      value replaces this one.
+
+      *Available since v1.3.0*.
+      """
     ]
   ]
 
@@ -1372,6 +1384,7 @@ defmodule Broadway do
 
   defp test_messages(broadway, data, batch_mode, opts) when is_broadway_name(broadway) do
     metadata = opts |> Keyword.fetch!(:metadata) |> Map.new()
+    batcher = Keyword.fetch!(opts, :batcher)
 
     acknowledger =
       Keyword.get(opts, :acknowledger, fn _data, ack_ref ->
@@ -1383,7 +1396,14 @@ defmodule Broadway do
     messages =
       Enum.map(data, fn data ->
         ack = acknowledger.(data, {self(), ref})
-        %Message{data: data, acknowledger: ack, batch_mode: batch_mode, metadata: metadata}
+
+        %Message{
+          data: data,
+          acknowledger: ack,
+          batch_mode: batch_mode,
+          metadata: metadata,
+          batcher: batcher
+        }
       end)
 
     :ok = push_messages(broadway, messages)

test/broadway_test.exs

@@ -791,6 +791,23 @@ defmodule BroadwayTest do
       assert_receive {:ack, ^ref, [%{status: :ok}, %{status: :ok}], []}
     end
 
+    test "messages use the batcher set through the :batcher option" do
+      broadway_name = new_unique_name()
+
+      {:ok, _broadway} =
+        Broadway.start_link(Forwarder,
+          name: broadway_name,
+          context: %{test_pid: self()},
+          producer: [module: {ManualProducer, []}],
+          processors: [default: [concurrency: 1, min_demand: 1, max_demand: 2]],
+          batchers: [default: [batch_size: 2], custom_batcher: [batch_size: 2]]
+        )
+
+      ref = Broadway.test_batch(broadway_name, [1, 2], batcher: :custom_batcher)
+      assert_receive {:batch_handled, :custom_batcher, [%{data: 1}, %{data: 2}]}
+      assert_receive {:ack, ^ref, [%{data: 1, status: :ok}, %{data: 2, status: :ok}], []}
+    end
+
     test "failed messages are marked as {:failed, reason}", %{broadway: broadway} do
       ref = Broadway.test_message(broadway, :fail)
       assert_receive {:ack, ^ref, _, [%{status: {:failed, "Failed message"}}]}