diff --git a/README.md b/README.md
index f16ed6e..ad2cb72 100644
--- a/README.md
+++ b/README.md
@@ -72,11 +72,8 @@ defmodule Products.Schema do
     directive :key, fields: "id"
 
     # Any subgraph contributing fields MUST define a _resolve_reference field.
-    # Note that implementing the reference resolver with function capture does not work at the moment. Hence, the examples below use an anonymous function.
     field :_resolve_reference, :product do
-      resolve(fn %{__typename: "Product", id: id} = entity, _info ->
-        {:ok, Map.merge(entity, %{name: "ACME Anvil", price: 10000})}
-      end)
+      resolve &Products.find_by_id/2
     end
 
     field :id, non_null(:id)
@@ -238,6 +235,60 @@ defmodule Example.Schema do
 end
 ```
 
+### Using Dataloader in \_resolve_reference queries
+
+You can use Dataloader in to resolve references to specific objects, but it requires manually setting up the batch and item key, as the field has no parent. Resolution for both \_resolve\_reference fields are functionally equivalent.
+
+```elixir
+defmodule Example.Schema do
+  use Absinthe.Schema
+  use Absinthe.Federation.Schema
+
+  import Absinthe.Resolution.Helpers, only: [on_load: 2, dataloader: 2]
+
+  def context(ctx) do
+    loader =
+      Dataloader.new()
+      |> Dataloader.add_source(Example.Loader, Dataloader.Ecto.new(Example.Repo))
+
+    Map.put(ctx, :loader, loader)
+  end
+
+  def plugins do
+    [Absinthe.Middleware.Dataloader] ++ Absinthe.Plugin.defaults()
+  end
+
+  object :item do
+    key_fields("item_id")
+
+    # Using the dataloader/2 resolution helper
+    field :_resolve_reference, :item do
+      resolve dataloader(Example.Loader, fn _parent, args, _res ->
+                %{batch: {{:one, Example.Item}, %{}}, item: [item_id: args.item_id]}
+              end)
+    end
+  end
+
+  object :verbose_item do
+    key_fields("item_id")
+
+    # Using the on_load/2 resolution helper
+    field :_resolve_reference, :verbose_item do
+      resolve fn %{item_id: id}, %{context: %{loader: loader}} ->
+        batch_key = {:one, Example.Item, %{}}
+        item_key = [item_id: id]
+
+        loader
+        |> Dataloader.load(Example.Loader, batch_key, item_key)
+        |> on_load(fn loader ->
+          result = Dataloader.get(loader, Example.Loader, batch_key, item_key)
+          {:ok, result}
+        end)
+    end
+  end
+end
+```
+
 ### Resolving structs in \_entities queries
 
 If you need to resolve your struct to a specific type in your schema you can implement the `Absinthe.Federation.Schema.EntityUnion.Resolver` protocol like this:
diff --git a/lib/absinthe/federation/schema/entities_field.ex b/lib/absinthe/federation/schema/entities_field.ex
index b49c581..52d6982 100644
--- a/lib/absinthe/federation/schema/entities_field.ex
+++ b/lib/absinthe/federation/schema/entities_field.ex
@@ -83,121 +83,153 @@ defmodule Absinthe.Federation.Schema.EntitiesField do
   end
 
   def call(%{state: :unresolved} = resolution, _args) do
-    resolutions = resolver(resolution.source, resolution.arguments, resolution)
-
-    resolvers =
-      resolutions
-      |> Enum.group_by(fn %{middleware: [middleware | _remaining_middleware]} = r ->
-        case middleware do
-          {Absinthe.Middleware.Dataloader, {loader, _fun}} ->
-            {source, _} = find_relevant_dataloader(loader)
-            {:dataloader, source}
-
-          _ ->
-            {:resolver, r}
-        end
-      end)
-      |> Enum.flat_map(fn
-        {{:dataloader, _}, v} = _resolvers -> Enum.take(v, 1)
-        {{:resolver, _}, v} = _resolvers -> v
-      end)
+    resolution_acc = resolution_accumulator(resolution)
+
+    # Run pre-resolution plugins, such as async/batch and dataloader.
+    resolution_acc = run_callbacks(resolution_acc.schema.plugins(), :before_resolution, resolution_acc)
+
+    representations_to_resolve =
+      Enum.map(resolution.arguments.representations, &resolve_reference_field(&1, resolution_acc))
+
+    # Resolve representations first time
+    resolution_acc = Enum.reduce(representations_to_resolve, resolution_acc, &resolve_field_representation/2)
+
+    # If any representation fields are suspended (i.e async or dataloaded),
+    # run the plugins and resolve_representation pipeline again.
+    resolution_acc =
+      if Enum.any?(resolution_acc.path, &(&1.state == :suspended)) do
+        representations_to_resolve = resolution_acc.path
+
+        resolution_acc =
+          run_callbacks(resolution_acc.schema.plugins(), :before_resolution, resolution_acc)
+          |> Map.put(:path, [])
+
+        Enum.reduce(representations_to_resolve, resolution_acc, &resolve_field_representation/2)
+      else
+        resolution_acc
+      end
+
+    # Run post-resolution plugins
+    resolution_acc = run_callbacks(resolution_acc.schema.plugins(), :after_resolution, resolution_acc)
+    representations = resolution_acc.path
 
-    value =
-      resolvers
-      |> Enum.map(&reduce_resolution/1)
-      |> List.flatten()
-      |> Map.new()
-
-    res =
-      resolution.arguments.representations
-      |> Enum.reduce(%{errors: [], value: []}, fn r, acc ->
-        case Map.get(value, r) do
-          {:error, err} ->
-            Map.update!(acc, :errors, &[err | &1])
-
-          {result, errors} ->
-            acc |> Map.update!(:value, &[result | &1]) |> Map.update!(:errors, &(errors ++ &1))
-
-          result ->
-            Map.update!(acc, :value, &[result | &1])
-        end
+    # Collect values and errors
+    resolution_acc =
+      Enum.reduce(representations, resolution_acc, fn representation, acc ->
+        representation_errors = representation.errors
+        representation_value = representation.value
+
+        acc
+        |> Map.update!(:value, fn value -> [representation_value | value] end)
+        |> Map.update!(:errors, fn errors -> representation_errors ++ errors end)
       end)
+      |> then(&reverse_values_and_errors/1)
+      |> then(&set_final_state/1)
 
-    %{
-      resolution
-      | state: :resolved,
-        errors: Enum.reverse(res[:errors]),
-        value: Enum.reverse(res[:value])
-    }
+    %{resolution | state: resolution_acc.state, value: resolution_acc.value, errors: resolution_acc.errors}
   end
 
-  def call(res, _args), do: res
-
-  def resolver(parent, %{representations: representations}, resolution) do
-    Enum.map(representations, &entity_accumulator(&1, parent, resolution))
+  defp reverse_values_and_errors(res) do
+    res
+    |> Map.update!(:value, &Enum.reverse/1)
+    |> Map.update!(:errors, &Enum.reverse/1)
   end
 
-  defp entity_accumulator(representation, parent, %{schema: schema} = resolution) do
-    typename = Map.get(representation, "__typename")
+  defp set_final_state(res) do
+    if Enum.all?(res.path, &(&1.state == :resolved)) do
+      Map.put(res, :state, :resolved)
+    else
+      paths = Enum.map(res.path, &%{state: &1.state, value: &1.value, errors: &1.errors})
+      raise "Expected all fields to be resolved, but got: #{paths}"
+    end
+  end
 
-    fun =
-      schema
-      |> Absinthe.Schema.lookup_type(typename)
-      |> resolve_representation(parent, representation, resolution)
+  # These are the fields to be threaded through every single representation resolution.
+  defp resolution_accumulator(resolution) do
+    %{resolution | path: [], errors: [], value: []}
+  end
 
-    resolution = Map.put(resolution, :arguments, Map.put(resolution.arguments, :representation, representation))
+  # Resolve a single representation and accumulate it to field resolution
+  # under path key. If a field is already resolved, do not run any middlewares.
+  defp resolve_field_representation(representation, acc) do
+    if representation.state == :resolved do
+      new_path = acc.path ++ [representation]
 
-    Absinthe.Resolution.call(resolution, fun)
+      acc
+      |> Map.put(:path, new_path)
+    else
+      local_res =
+        representation
+        |> Map.put(:context, acc.context)
+        |> Map.put(:acc, acc.acc)
+
+      result = reduce_resolution(local_res)
+      new_path = acc.path ++ [result]
+
+      acc
+      |> Map.put(:context, result.context)
+      |> Map.put(:acc, result.acc)
+      |> Map.put(:path, new_path)
+    end
   end
 
-  defp resolve_representation(
-         %struct_type{fields: fields},
-         parent,
-         representation,
-         resolution
-       )
-       when struct_type in [Absinthe.Type.Object, Absinthe.Type.Interface],
-       do: resolve_reference(fields[:_resolve_reference], parent, representation, resolution)
+  defp resolve_reference_field(representation, resolution_acc) do
+    typename = Map.get(representation, "__typename")
 
-  defp resolve_representation(_schema_type, _parent, representation, _schema),
-    do:
-      {:error,
-       "The _entities resolver tried to load an entity for type '#{Map.get(representation, "__typename")}', but no object type of that name was found in the schema"}
+    %Absinthe.Resolution{
+      schema: schema,
+      source: source,
+      context: context,
+      adapter: adapter,
+      definition: definition,
+      parent_type: parent_type
+    } = resolution_acc
 
-  defp resolve_reference(nil, _parent, representation, %{context: context} = _resolution) do
     args = convert_keys_to_atom(representation, context)
 
-    fn _, _ -> {:ok, args} end
+    field =
+      schema
+      |> Absinthe.Schema.lookup_type(typename)
+      |> resolve_representation(representation)
+
+    %Absinthe.Resolution{
+      arguments: args,
+      schema: schema,
+      definition: definition,
+      parent_type: parent_type,
+      source: source,
+      state: :unresolved,
+      value: nil,
+      errors: [],
+      middleware: field.middleware,
+      context: context,
+      adapter: adapter
+    }
   end
 
-  defp resolve_reference(
-         %{middleware: middleware},
-         parent,
-         representation,
-         %{schema: schema, context: context} = resolution
-       ) do
-    args = convert_keys_to_atom(representation, context)
-
-    middleware
-    |> maybe_unshim(schema)
-    |> Enum.find(nil, &only_resolver_middleware/1)
-    |> case do
-      {_, resolve_ref_func} when is_function(resolve_ref_func, 2) ->
-        fn _, _ -> resolve_ref_func.(args, resolution) end
+  defp resolve_representation(%struct_type{fields: fields}, _representation)
+       when struct_type in [Absinthe.Type.Object, Absinthe.Type.Interface] do
+    resolve_reference(fields[:_resolve_reference])
+  end
 
-      {_, resolve_ref_func} when is_function(resolve_ref_func, 3) ->
-        fn _, _ -> resolve_ref_func.(parent, args, resolution) end
+  defp resolve_representation(_schema_type, representation),
+    do:
+      {:error,
+       "The _entities resolver tried to load an entity for type '#{Map.get(representation, "__typename")}', but no object type of that name was found in the schema"}
 
-      _ ->
-        fn _, _ -> {:ok, args} end
+  defp resolve_reference(field) do
+    # When there is a field _resolve_reference, set it up so the resolution pipeline can be run
+    # on it.
+    if field do
+      field
+    else
+      # When there is no field name _resolve_reference defined on the key object, create
+      # a stub middleware that returns arguments as the field resolution.
+      middleware = {{Absinthe.Resolution, :call}, fn args, _res -> {:ok, args} end}
+      %Absinthe.Resolution{middleware: [middleware], state: :unresolved, value: nil}
     end
   end
 
-  defp maybe_unshim([{{Absinthe.Middleware, :shim}, {_, _, _}}] = middleware, schema),
-    do: Absinthe.Middleware.unshim(middleware, schema)
-
-  defp maybe_unshim(middleware, _schema), do: middleware
-
   defp convert_keys_to_atom(map, context) when is_map(map) do
     Map.new(map, fn {k, v} ->
       k = convert_key(k, context)
@@ -219,13 +251,41 @@ defmodule Absinthe.Federation.Schema.EntitiesField do
     |> String.to_atom()
   end
 
-  defp adapter_has_to_internal_name_modifier?(adapter) do
-    Keyword.get(adapter.__info__(:functions), :to_internal_name) == 2
+  defp run_callbacks(plugins, callback, acc) do
+    Enum.reduce(plugins, acc, &apply(&1, callback, [&2]))
+  end
+
+  defp reduce_resolution(%Absinthe.Resolution{middleware: []} = res), do: res
+
+  defp reduce_resolution(%Absinthe.Resolution{middleware: [middleware | remaining_middleware]} = res) do
+    case call_middleware(middleware, %{res | middleware: remaining_middleware}) do
+      %{state: :suspended} = res ->
+        res
+
+      res ->
+        reduce_resolution(res)
+    end
   end
 
-  defp only_resolver_middleware({{Absinthe.Resolution, :call}, _}), do: true
+  defp call_middleware({{mod, fun}, opts}, res) do
+    apply(mod, fun, [res, opts])
+  end
 
-  defp only_resolver_middleware(_), do: false
+  defp call_middleware({mod, opts}, res) do
+    apply(mod, :call, [res, opts])
+  end
+
+  defp call_middleware(mod, res) when is_atom(mod) do
+    apply(mod, :call, [res, []])
+  end
+
+  defp call_middleware(fun, res) when is_function(fun, 2) do
+    fun.(res, [])
+  end
+
+  defp adapter_has_to_internal_name_modifier?(adapter) do
+    Keyword.get(adapter.__info__(:functions), :to_internal_name) == 2
+  end
 
   defp build_arguments(), do: [build_argument()]
 
@@ -246,68 +306,4 @@ defmodule Absinthe.Federation.Schema.EntitiesField do
         }
       }
     }
-
-  defp reduce_resolution(%{state: :resolved} = res) do
-    {res.arguments.representation, {res.value, res.errors}}
-  end
-
-  defp reduce_resolution(%{middleware: []} = res), do: res
-
-  defp reduce_resolution(%{middleware: [middleware | remaining_middleware]} = res) do
-    call_middleware(middleware, %{res | middleware: remaining_middleware})
-  end
-
-  defp call_middleware({Absinthe.Middleware.Dataloader, {loader, _fun}}, %{
-         arguments: %{representations: args},
-         schema: schema
-       }) do
-    {source, typename} = find_relevant_dataloader(loader)
-
-    key_field =
-      Absinthe.Schema.lookup_type(schema, typename)
-      |> Absinthe.Type.meta()
-      |> Map.get(:key_fields, "id")
-
-    representations =
-      args
-      |> Enum.reject(fn arg ->
-        Map.get(arg, "__typename") != typename
-      end)
-
-    ids =
-      representations
-      |> Enum.map(fn arg -> Map.get(arg, key_field) end)
-
-    loader
-    |> Dataloader.load_many(source, %{__typename: typename}, ids)
-    |> Dataloader.run()
-    |> Dataloader.get_many(source, %{__typename: typename}, ids)
-    |> Enum.zip(representations)
-    |> Enum.map(fn {res, arg} ->
-      case res do
-        {:ok, data} -> {arg, data}
-        {:error, _} = e -> {arg, e}
-        data -> {arg, data}
-      end
-    end)
-  end
-
-  defp call_middleware({_mod, {fun, args}}, resolution) do
-    with {:ok, res} <- fun.(args) do
-      {resolution.arguments.representation, res}
-    else
-      err -> {resolution.arguments.representation, err}
-    end
-  end
-
-  defp find_relevant_dataloader(%Dataloader{sources: sources}) do
-    {source, loader} =
-      Enum.find(sources, fn {_, source} ->
-        Dataloader.Source.pending_batches?(source)
-      end)
-
-    %{batches: batches} = loader
-    {{:_entities, %{__typename: typename}}, _} = Enum.at(batches, 0)
-    {source, typename}
-  end
 end
diff --git a/test/absinthe/federation/schema/entities_field/dataloader_test.exs b/test/absinthe/federation/schema/entities_field/dataloader_test.exs
new file mode 100644
index 0000000..505a6f0
--- /dev/null
+++ b/test/absinthe/federation/schema/entities_field/dataloader_test.exs
@@ -0,0 +1,129 @@
+defmodule Absinthe.Federation.Schema.EntitiesField.DataloaderTest do
+  use Absinthe.Federation.Case, async: true
+
+  setup do
+    {:ok, source} = start_supervised(Example.Source)
+    Example.Source.put(%{"1" => %Example.Item{item_id: "1"}, "3" => %Example.Item{item_id: "3"}})
+
+    %{source: source}
+  end
+
+  describe "resolver with dataloader" do
+    defmodule ExampleDataloaderSchema do
+      use Absinthe.Schema
+      use Absinthe.Federation.Schema
+
+      import Absinthe.Resolution.Helpers, only: [on_load: 2, dataloader: 3]
+
+      def context(ctx) do
+        loader =
+          Dataloader.new()
+          |> Dataloader.add_source(Example.Source, Dataloader.KV.new(&Example.Source.run_batch/2))
+
+        Map.put(ctx, :loader, loader)
+      end
+
+      def plugins do
+        [Absinthe.Middleware.Dataloader] ++ Absinthe.Plugin.defaults()
+      end
+
+      query do
+      end
+
+      object :normal_item do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :normal_item do
+          resolve fn %{item_id: id, __typename: typename}, _res ->
+            {:ok, %{item_id: id, __typename: typename}}
+          end
+        end
+      end
+
+      object :dataloaded_item do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :dataloaded_item do
+          resolve dataloader(
+                    Example.Source,
+                    fn _parent, args, _res -> %{batch: {:one, Example.Item, %{}}, item: args.item_id} end,
+                    callback: fn item, _parent, _args ->
+                      if item do
+                        item = Map.drop(item, [:__struct__])
+                        item = Map.put(item, :__typename, "DataloadedItem")
+                        {:ok, item}
+                      else
+                        {:ok, item}
+                      end
+                    end
+                  )
+        end
+      end
+
+      object :on_load_item do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :on_load_item do
+          resolve fn %{item_id: id}, %{context: %{loader: loader}} ->
+            batch_key = {:one, Example.Item, %{}}
+            item_key = id
+
+            loader
+            |> Dataloader.load(Example.Source, batch_key, item_key)
+            |> on_load(fn loader ->
+              result = Dataloader.get(loader, Example.Source, batch_key, item_key)
+
+              if result do
+                result = Map.drop(result, [:__struct__])
+                {:ok, Map.put(result, :__typename, "DataloadedItem")}
+              else
+                {:ok, nil}
+              end
+            end)
+          end
+        end
+      end
+    end
+
+    test "handles dataloader resolvers" do
+      query = """
+        query {
+          _entities(representations: [
+            {
+              __typename: "DataloadedItem",
+              item_id: "1"
+            },
+            {
+              __typename: "NormalItem",
+              item_id: "1"
+            },
+            {
+              __typename: "OnLoadItem",
+              item_id: "1"
+            },
+            {
+              __typename: "OnLoadItem",
+              item_id: "2"
+            }
+          ]) {
+            ...on DataloadedItem {
+              item_id
+            }
+            ...on NormalItem {
+              item_id
+            }
+            ...on OnLoadItem {
+              item_id
+            }
+          }
+        }
+      """
+
+      assert {:ok, %{data: %{"_entities" => [%{"item_id" => "1"}, %{"item_id" => "1"}, %{"item_id" => "1"}, nil]}}} =
+               Absinthe.run(query, ExampleDataloaderSchema, variables: %{})
+    end
+  end
+end
diff --git a/test/absinthe/federation/schema/entities_field/middleware_test.exs b/test/absinthe/federation/schema/entities_field/middleware_test.exs
new file mode 100644
index 0000000..29183ba
--- /dev/null
+++ b/test/absinthe/federation/schema/entities_field/middleware_test.exs
@@ -0,0 +1,116 @@
+defmodule Absinthe.Federation.Schema.EntitiesField.MiddlewareTest do
+  use Absinthe.Federation.Case, async: true
+
+  describe "resolve a function capture" do
+    defmodule FunctionCaptureSchema do
+      use Absinthe.Schema
+      use Absinthe.Federation.Schema
+
+      query do
+      end
+
+      @impl Absinthe.Schema
+      def middleware(middleware, _field, %{identifier: :item_with_module_middleware}) do
+        middleware ++ [Example.Middleware]
+      end
+
+      def middleware(middleware, _field, _object) do
+        middleware
+      end
+
+      object :item_with_module_middleware do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :item_with_module_middleware
+      end
+
+      object :item_with_function_middleware do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :item_with_function_middleware do
+          resolve &__MODULE__.get_item/2
+
+          middleware fn res, _ ->
+            value = Map.update!(res.value, :item_id, &"FunctionMiddleware:#{&1}")
+            Map.put(res, :value, value)
+          end
+        end
+      end
+
+      object :item_with_function_capture do
+        key_fields("item_id")
+        field :item_id, :string
+
+        field :_resolve_reference, :item_with_function_capture do
+          resolve &__MODULE__.get_item/2
+        end
+      end
+
+      def get_item(args, _res) do
+        {:ok, args}
+      end
+    end
+
+    test "handles a post-resolution middleware" do
+      query = """
+        query {
+          _entities(representations: [
+            {
+              __typename: "ItemWithFunctionMiddleware",
+              item_id: "1"
+            }
+          ]) {
+            ...on ItemWithFunctionMiddleware {
+              item_id
+            }
+          }
+        }
+      """
+
+      assert {:ok, %{data: %{"_entities" => [%{"item_id" => "FunctionMiddleware:1"}]}}} =
+               Absinthe.run(query, FunctionCaptureSchema, variables: %{})
+    end
+
+    test "handles a module-based middleware" do
+      query = """
+        query {
+          _entities(representations: [
+            {
+              __typename: "ItemWithModuleMiddleware",
+              item_id: "1"
+            }
+          ]) {
+            ...on ItemWithModuleMiddleware {
+              item_id
+            }
+          }
+        }
+      """
+
+      assert {:ok, %{data: %{"_entities" => [%{"item_id" => "ModuleMiddleware:1"}]}}} =
+               Absinthe.run(query, FunctionCaptureSchema, variables: %{})
+    end
+
+    test "handles a function capture" do
+      query = """
+        query {
+          _entities(representations: [
+            {
+              __typename: "ItemWithFunctionCapture",
+              item_id: "1"
+            }
+          ]) {
+            ...on ItemWithFunctionCapture {
+              item_id
+            }
+          }
+        }
+      """
+
+      assert {:ok, %{data: %{"_entities" => [%{"item_id" => "1"}]}}} =
+               Absinthe.run(query, FunctionCaptureSchema, variables: %{})
+    end
+  end
+end
diff --git a/test/absinthe/federation/schema/entities_field_test.exs b/test/absinthe/federation/schema/entities_field_test.exs
index 2cfa2bd..58ef4d7 100644
--- a/test/absinthe/federation/schema/entities_field_test.exs
+++ b/test/absinthe/federation/schema/entities_field_test.exs
@@ -66,7 +66,7 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
 
         field :_resolve_reference, :product do
           resolve(fn _, %{upc: upc} = args, _ ->
-            async(fn _ ->
+            async(fn ->
               case upc do
                 "123" -> {:ok, args}
                 "456" -> {:ok, args}
@@ -285,7 +285,7 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
       use Absinthe.Schema
       use Absinthe.Federation.Schema
 
-      import Absinthe.Resolution.Helpers, only: [dataloader: 1]
+      import Absinthe.Resolution.Helpers, only: [dataloader: 2]
 
       def context(ctx) do
         loader =
@@ -312,9 +312,9 @@ defmodule Absinthe.Federation.Schema.EntitiesFieldTest do
         field :item_id, :string
 
         field :_resolve_reference, :spec_item do
-          resolve(fn _root, %{item_id: id} = args, info ->
-            dataloader(SpecItem.Loader).(id, args, info)
-          end)
+          resolve dataloader(SpecItem.Loader, fn _parent, args, _res ->
+                    %{batch: {{:one, SpecItem}, %{}}, item: args.item_id}
+                  end)
         end
       end
     end
diff --git a/test/support/example_item.ex b/test/support/example_item.ex
new file mode 100644
index 0000000..601df2e
--- /dev/null
+++ b/test/support/example_item.ex
@@ -0,0 +1,3 @@
+defmodule Example.Item do
+  defstruct [:item_id]
+end
diff --git a/test/support/example_middleware.ex b/test/support/example_middleware.ex
new file mode 100644
index 0000000..367fac6
--- /dev/null
+++ b/test/support/example_middleware.ex
@@ -0,0 +1,18 @@
+defmodule Example.Middleware do
+  @behaviour Absinthe.Middleware
+
+  @impl true
+  def call(%Absinthe.Resolution{} = res, _opts) do
+    case res do
+      %{arguments: %{item_id: item_id, __typename: "ItemWithModuleMiddleware"}} ->
+        value = %{item_id: "ModuleMiddleware:#{item_id}", __typename: "ItemWithModuleMiddleware"}
+
+        res
+        |> Map.put(:value, value)
+        |> Map.put(:state, :resolved)
+
+      res ->
+        res
+    end
+  end
+end
diff --git a/test/support/example_source.ex b/test/support/example_source.ex
new file mode 100644
index 0000000..85898a5
--- /dev/null
+++ b/test/support/example_source.ex
@@ -0,0 +1,23 @@
+defmodule Example.Source do
+  use Agent
+
+  def start_link(initial_value) do
+    Agent.start_link(fn -> initial_value end, name: __MODULE__)
+  end
+
+  def value do
+    Agent.get(__MODULE__, & &1)
+  end
+
+  def put(value) do
+    Agent.update(__MODULE__, fn _ -> value end)
+  end
+
+  def run_batch({:one, Example.Item, %{}}, items) do
+    value = Agent.get(__MODULE__, & &1)
+
+    for item <- items, into: %{} do
+      {item, Map.get(value, item)}
+    end
+  end
+end