Repository: plataformatec/broadway (elixir)

Issue:
Broadway.test_messages/2 doesn't send ack when using batchers
Here's a test that reproduces it:

```elixir
defmodule EmptyProducer do
  use GenStage
  @behaviour Broadway.Producer

  @impl true
  def init(_args) do
    {:producer, []}
  end

  @impl true
  def handle_demand(_demand, state) do
    {:noreply, [], state}
  end
end

defmodule MyBroadway do
  use Broadway

  @impl true
  def handle_message(:default, message, _context) do
    message
  end
end

defmodule BugTest do
  use ExUnit.Case

  test "works", c do
    {:ok, pid} =
      Broadway.start_link(MyBroadway,
        name: c.test,
        producers: [
          default: [
            module: {EmptyProducer, []}
          ]
        ],
        processors: [
          default: [
            stages: 1
          ]
        ]
      )

    ref = Broadway.test_messages(pid, [1])
    assert_receive {:ack, ^ref, [%{data: 1}], []}
  end

  test "works with batchers", c do
    {:ok, pid} =
      Broadway.start_link(MyBroadway,
        name: c.test,
        producers: [
          default: [
            module: {EmptyProducer, []}
          ]
        ],
        batchers: [
          default: [
            stages: 1
          ]
        ],
        processors: [
          default: [
            stages: 1
          ]
        ]
      )

    ref = Broadway.test_messages(pid, [1])
    assert_receive {:ack, ^ref, [%{data: 1}], []}
  end
end
```

The 2nd test fails with:

```
  1) test works with batchers (BugTest)
     test/bug_test.exs:48
     No message matching {:ack, ^ref, [%{data: 1}], []} after 100ms.
     The following variables were pinned:
       ref = #Reference<0.1852397191.3371696133.32665>
     The process mailbox is empty.
     code: assert_receive {:ack, ^ref, [%{data: 1}], []}
     stacktrace:
       test/bug_test.exs:70: (test)
```

Pull request description:
Add `:batch_mode` and default to `:flush` for test messages
Closes #92.

Worth mentioning that for people using default batch size of 100 the following works:

```elixir
ref = Broadway.test_messages(broadway, [1, 2, 3, 4])
assert_receive {:ack, ^ref, [%{data: 1}, %{data: 2}, %{data: 3}, %{data: 4}], []}
```

so it's not that people would necessarily get a lot more acknowledgments, we only got them (which required us to tweak our tests) because we had low batch size in tests.

Patch:
diff --git a/lib/broadway.ex b/lib/broadway.ex
index 4c58769..0532bf2 100644
--- a/lib/broadway.ex
+++ b/lib/broadway.ex
@@ -297,13 +297,20 @@ defmodule Broadway do
       # Now assert the database side-effects
       ...
 
-  Keep in mind that multiple acknowledgement messages may be sent.
-  For example, if the batcher in the example above has size of 2,
-  then two batches would be created and therefore two ack messages
-  would be sent. Similarly, if any of the messages fail when
-  processed, an acknowledgement of their failure may be sent early
-  on. On the positive side, if you always push just a single test
-  message, then there is always one acknowledgment.
+  When testing pipelines that are using batchers there are additional
+  considerations. By default, the batch is only delivered when either
+  its size or its timeout has been reached, but that is often impractical
+  for testing, you may not necessarily want to send a lot of data or
+  wait a lot of time for the batch to flush. For this reason, when using
+  `test_messages/2`, by default the messages have their `:batch_mode`
+  set to `:flush` which means the batch will be immediately delivered.
+  Depending on the batch size and the batching mode you might get multiple
+  acknowledgment messages. For example, if the batcher in the example above
+  has size of 2, each event when batched would be immediately flushed and
+  so you would get a total of 3 acknowledment messages, one for each event.
+  Similarly, if any of the messages fail when processed, an acknowledgement
+  of their failure may be sent early on. On the positive side, if you always
+  push just a single test message, then there is always one acknowledgment.
   """
 
   alias Broadway.{BatchInfo, Message, Options, Server, Producer}
@@ -533,6 +540,12 @@ defmodule Broadway do
 
   See "Testing" section in module documentation for more information.
 
+  ## Options
+
+  * `:batch_mode` - 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. Defaults to `:flush`.
+
   ## Examples
 
   For example, in your tests, you may do:
@@ -544,10 +557,13 @@ defmodule Broadway do
 
   """
   @spec test_messages(GenServer.server(), data :: [term]) :: reference
-  def test_messages(broadway, data) when is_list(data) do
+  def test_messages(broadway, data, opts \\ []) when is_list(data) and is_list(opts) do
+    batch_mode = Keyword.get(opts, :batch_mode, :flush)
     ref = make_ref()
     ack = {Broadway.CallerAcknowledger, {self(), ref}, :ok}
-    messages = Enum.map(data, &%Message{data: &1, acknowledger: ack})
+
+    messages = Enum.map(data, &%Message{data: &1, acknowledger: ack, batch_mode: batch_mode})
+
     :ok = push_messages(broadway, messages)
     ref
   end
diff --git a/lib/broadway/batcher.ex b/lib/broadway/batcher.ex
index 311889b..0e80d73 100644
--- a/lib/broadway/batcher.ex
+++ b/lib/broadway/batcher.ex
@@ -81,29 +81,41 @@ defmodule Broadway.Batcher do
 
   defp handle_events_per_batch_key([%{batch_key: batch_key} | _] = events, acc, state) do
     {current, pending_count, timer} = init_or_get_batch(batch_key, state)
-    {current, pending_count, events} = split_counting(batch_key, events, pending_count, current)
 
-    acc = deliver_or_update_batch(batch_key, current, pending_count, timer, acc, state)
+    {current, pending_count, events, flush?} =
+      split_counting(batch_key, events, pending_count, false, current)
+
+    acc = deliver_or_update_batch(batch_key, current, pending_count, flush?, timer, acc, state)
     handle_events_per_batch_key(events, acc, state)
   end
 
-  defp split_counting(batch_key, [%{batch_key: batch_key} = event | events], count, acc)
-       when count > 0,
-       do: split_counting(batch_key, events, count - 1, [event | acc])
+  defp split_counting(batch_key, [%{batch_key: batch_key} = event | events], count, flush?, acc)
+       when count > 0 do
+    flush? = flush? or event.batch_mode == :flush
+    split_counting(batch_key, events, count - 1, flush?, [event | acc])
+  end
 
-  defp split_counting(_batch_key, events, count, acc), do: {acc, count, events}
+  defp split_counting(_batch_key, events, count, flush?, acc), do: {acc, count, events, flush?}
 
-  defp deliver_or_update_batch(batch_key, current, 0, timer, acc, state) do
-    delete_batch(batch_key)
-    cancel_batch_timeout(timer)
-    [wrap_for_delivery(batch_key, current, state) | acc]
+  defp deliver_or_update_batch(batch_key, current, _pending_count, true, timer, acc, state) do
+    deliver_batch(batch_key, current, timer, acc, state)
+  end
+
+  defp deliver_or_update_batch(batch_key, current, 0, _flush?, timer, acc, state) do
+    deliver_batch(batch_key, current, timer, acc, state)
   end
 
-  defp deliver_or_update_batch(batch_key, current, pending_count, timer, acc, _state) do
+  defp deliver_or_update_batch(batch_key, current, pending_count, _flush?, timer, acc, _state) do
     put_batch(batch_key, {current, pending_count, timer})
     acc
   end
 
+  defp deliver_batch(batch_key, current, timer, acc, state) do
+    delete_batch(batch_key)
+    cancel_batch_timeout(timer)
+    [wrap_for_delivery(batch_key, current, state) | acc]
+  end
+
   ## General batch handling
 
   defp init_or_get_batch(batch_key, state) do
diff --git a/lib/broadway/message.ex b/lib/broadway/message.ex
index f311ae9..c60b194 100644
--- a/lib/broadway/message.ex
+++ b/lib/broadway/message.ex
@@ -20,6 +20,7 @@ defmodule Broadway.Message do
           acknowledger: {module, ack_ref :: term, data :: term},
           batcher: atom,
           batch_key: term,
+          batch_mode: :bulk | :flush,
           status: :ok | {:failed, reason :: binary}
         }
 
@@ -29,6 +30,7 @@ defmodule Broadway.Message do
             acknowledger: nil,
             batcher: :default,
             batch_key: :default,
+            batch_mode: :bulk,
             status: :ok
 
   @doc """
@@ -61,6 +63,23 @@ defmodule Broadway.Message do
     %Message{message | batch_key: batch_key}
   end
 
+  @doc """
+  Sets the batching mode for the message.
+
+  When the mode is `:bulk`, the batch that the message is in is delivered after
+  the batch size or the batch timeout is reached.
+
+  When the mode is `:flush`, the batch that the message is in is delivered
+  immediately after processing. Note it doesn't mean the batch contains only a single element
+  but rather that all messages receives from the processor are delivered without waiting.
+
+  The default mode for messages is `:bulk`.
+  """
+  @spec put_batch_mode(message :: Message.t(), mode :: :bulk | :flush) :: Message.t()
+  def put_batch_mode(%Message{} = message, mode) when mode in [:bulk, :flush] do
+    %Message{message | batch_mode: mode}
+  end
+
   @doc """
   Mark a message as failed.
 


Test patch:
diff --git a/test/broadway_test.exs b/test/broadway_test.exs
index 88e546c..87c19a9 100644
--- a/test/broadway_test.exs
+++ b/test/broadway_test.exs
@@ -353,7 +353,7 @@ defmodule BroadwayTest do
     end
 
     test "successful messages are marked as :ok", %{broadway: broadway} do
-      ref = Broadway.test_messages(broadway, [1, 2])
+      ref = Broadway.test_messages(broadway, [1, 2], batch_mode: :bulk)
       assert_receive {:batch_handled, [%{status: :ok}, %{status: :ok}]}
       assert_receive {:ack, ^ref, [%{status: :ok}, %{status: :ok}], []}
     end
@@ -364,7 +364,7 @@ defmodule BroadwayTest do
     end
 
     test "failed messages are not forwarded to the batcher", %{broadway: broadway} do
-      ref = Broadway.test_messages(broadway, [1, :fail, :fail, 4])
+      ref = Broadway.test_messages(broadway, [1, :fail, :fail, 4], batch_mode: :bulk)
 
       assert_receive {:batch_handled, [%{data: 1}, %{data: 4}]}
       assert_receive {:ack, ^ref, [%{status: :ok}, %{status: :ok}], []}
@@ -375,7 +375,7 @@ defmodule BroadwayTest do
     test "messages are marked as {:failed, reason} when an error is raised while processing",
          %{broadway: broadway, processor: processor} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, :raise, 4])
+               ref = Broadway.test_messages(broadway, [1, :raise, 4], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [%{status: {:failed, "due to an unhandled error"}}]}
@@ -389,7 +389,7 @@ defmodule BroadwayTest do
     test "messages are marked as {:failed, reason} on bad return",
          %{broadway: broadway, processor: processor} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, :bad_return, 4])
+               ref = Broadway.test_messages(broadway, [1, :bad_return, 4], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [%{status: {:failed, "due to an unhandled error"}}]}
@@ -404,7 +404,7 @@ defmodule BroadwayTest do
     test "messages are marked as {:failed, reason} on bad batcher",
          %{broadway: broadway, processor: processor} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, :bad_batcher, 4])
+               ref = Broadway.test_messages(broadway, [1, :bad_batcher, 4], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [%{status: {:failed, "due to an unhandled error"}}]}
@@ -431,7 +431,7 @@ defmodule BroadwayTest do
 
     test "processors do not crash on broken link", %{broadway: broadway, processor: processor} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, :broken_link, 4])
+               ref = Broadway.test_messages(broadway, [1, :broken_link, 4], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [], [%{status: {:failed, "due to an unhandled exit"}}]}
 
@@ -541,7 +541,7 @@ defmodule BroadwayTest do
     end
 
     test "generate batches based on :batch_size", %{broadway: broadway} do
-      Broadway.test_messages(broadway, Enum.to_list(1..40))
+      Broadway.test_messages(broadway, Enum.to_list(1..40), batch_mode: :bulk)
 
       assert_receive {:batch_handled, :odd, messages, %BatchInfo{batcher: :odd}}
                      when length(messages) == 10
@@ -608,7 +608,7 @@ defmodule BroadwayTest do
     end
 
     test "generate batches based on :batch_size", %{broadway: broadway} do
-      ref = Broadway.test_messages(broadway, Enum.to_list(1..10))
+      ref = Broadway.test_messages(broadway, Enum.to_list(1..10), batch_mode: :bulk)
 
       assert_receive {:ack, ^ref, [%{data: 1, batch_key: :odd}, %{data: 3, batch_key: :odd}], []}
 
@@ -628,6 +628,60 @@ defmodule BroadwayTest do
     end
   end
 
+  describe "put_batch_mode" do
+    setup do
+      test_pid = self()
+      broadway_name = new_unique_name()
+
+      handle_message = fn message, _ ->
+        message
+      end
+
+      context = %{
+        handle_message: handle_message,
+        handle_batch: fn batcher, batch, batch_info, _ ->
+          send(test_pid, {:batch_handled, batcher, batch_info})
+          batch
+        end
+      }
+
+      {:ok, broadway} =
+        Broadway.start_link(CustomHandlers,
+          name: broadway_name,
+          context: context,
+          producers: [default: [module: {ManualProducer, []}]],
+          processors: [default: []],
+          batchers: [default: [batch_size: 2, batch_timeout: 20]]
+        )
+
+      %{broadway: broadway}
+    end
+
+    test "flush by default", %{broadway: broadway} do
+      ref = Broadway.test_messages(broadway, [1, 2, 3, 4, 5])
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 1}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 2}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 3}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 4}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 5}], []}
+    end
+
+    test "bulk", %{broadway: broadway} do
+      ref = Broadway.test_messages(broadway, [1, 2, 3, 4, 5], batch_mode: :bulk)
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 1}, %{data: 2}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 3}, %{data: 4}], []}
+      assert_receive {:batch_handled, :default, _}
+      assert_receive {:ack, ^ref, [%{data: 5}], []}
+    end
+  end
+
   describe "transformer" do
     setup tags do
       broadway_name = new_unique_name()
@@ -832,7 +886,7 @@ defmodule BroadwayTest do
     end
 
     test "only the messages in the crashing processor are lost", %{broadway: broadway} do
-      Broadway.test_messages(broadway, [1, 2, :kill_processor, 3, 4, 5])
+      Broadway.test_messages(broadway, [1, 2, :kill_processor, 3, 4, 5], batch_mode: :bulk)
 
       assert_receive {:message_handled, %{data: 1}}
       assert_receive {:message_handled, %{data: 2}}
@@ -845,7 +899,7 @@ defmodule BroadwayTest do
     end
 
     test "batches are created normally (without the lost messages)", %{broadway: broadway} do
-      Broadway.test_messages(broadway, [1, 2, :kill_processor, 3, 4, 5])
+      Broadway.test_messages(broadway, [1, 2, :kill_processor, 3, 4, 5], batch_mode: :bulk)
 
       assert_receive {:batch_handled, messages}
       values = messages |> Enum.map(& &1.data)
@@ -904,15 +958,15 @@ defmodule BroadwayTest do
     end
 
     test "batcher will be restarted in order to handle other messages", %{broadway: broadway} do
-      Broadway.test_messages(broadway, [1, 2])
+      Broadway.test_messages(broadway, [1, 2], batch_mode: :bulk)
       assert_receive {:batch_handled, _, batcher1} when is_pid(batcher1)
       ref = Process.monitor(batcher1)
 
-      Broadway.test_messages(broadway, [:kill_batcher, 3])
+      Broadway.test_messages(broadway, [:kill_batcher, 3], batch_mode: :bulk)
       assert_receive {:batch_handled, _, ^batcher1}
       assert_receive {:DOWN, ^ref, _, _, _}
 
-      Broadway.test_messages(broadway, [4, 5])
+      Broadway.test_messages(broadway, [4, 5], batch_mode: :bulk)
       assert_receive {:batch_handled, _, batcher2} when is_pid(batcher2)
       assert batcher1 != batcher2
     end
@@ -939,7 +993,8 @@ defmodule BroadwayTest do
     end
 
     test "only the messages in the crashing batcher are lost", %{broadway: broadway} do
-      ref = Broadway.test_messages(broadway, [1, 2, :kill_batcher, 3, 4, 5, 6, 7])
+      ref =
+        Broadway.test_messages(broadway, [1, 2, :kill_batcher, 3, 4, 5, 6, 7], batch_mode: :bulk)
 
       assert_receive {:ack, ^ref, successful, _failed}
       values = Enum.map(successful, & &1.data)
@@ -1014,7 +1069,7 @@ defmodule BroadwayTest do
     test "all messages in the batch are marked as {:failed, reason} when an error is raised",
          %{broadway: broadway, consumer: consumer} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, 2, :raise, 3])
+               ref = Broadway.test_messages(broadway, [1, 2, :raise, 3], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [
@@ -1031,7 +1086,7 @@ defmodule BroadwayTest do
     test "all messages in the batch are marked as {:failed, reason} on bad return",
          %{broadway: broadway, consumer: consumer} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, 2, :bad_return, 3])
+               ref = Broadway.test_messages(broadway, [1, 2, :bad_return, 3], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [
@@ -1055,7 +1110,7 @@ defmodule BroadwayTest do
 
       assert capture_log(fn ->
                Broadway.push_messages(broadway, [one, two, raise, three])
-               ref = Broadway.test_messages(broadway, [1, 2, 3, 4])
+               ref = Broadway.test_messages(broadway, [1, 2, 3, 4], batch_mode: :bulk)
                assert_receive {:ack, ^ref, [%{data: 1}, %{data: 2}, %{data: 3}, %{data: 4}], []}
              end) =~ "[error] ** (UndefinedFunctionError) function Unknown.ack/3 is undefined"
 
@@ -1065,7 +1120,7 @@ defmodule BroadwayTest do
     test "consumers do not crash on broken link",
          %{broadway: broadway, consumer: consumer} do
       assert capture_log(fn ->
-               ref = Broadway.test_messages(broadway, [1, 2, :broken_link, 3])
+               ref = Broadway.test_messages(broadway, [1, 2, :broken_link, 3], batch_mode: :bulk)
 
                assert_receive {:ack, ^ref, [],
                                [
@@ -1142,7 +1197,7 @@ defmodule BroadwayTest do
 
     test "shutting down broadway waits until all events are processed even on incomplete batches",
          %{broadway: broadway} do
-      ref = Broadway.test_messages(broadway, [1, 2])
+      ref = Broadway.test_messages(broadway, [1, 2], batch_mode: :bulk)
       Process.exit(broadway, :shutdown)
       assert_receive {:ack, ^ref, [%{data: 1}, %{data: 2}], []}
     end