1 files changed, 7 insertions, 10 deletions
diff --git a/spec/support/helpers/graphql_helpers.rb b/spec/support/helpers/graphql_helpers.rb
index 5bb2cf9dd9e..1eaa7603af0 100644
--- a/spec/support/helpers/graphql_helpers.rb
+++ b/spec/support/helpers/graphql_helpers.rb
@@ -4,20 +4,17 @@ module GraphqlHelpers
     kls[name].call(obj, args, ctx)
   end
 
-  # Runs a block inside a GraphQL::Batch wrapper
+  # Runs a block inside a BatchLoader::Executor wrapper
   def batch(max_queries: nil, &blk)
     wrapper = proc do
-      GraphQL::Batch.batch do
-        result = yield
-
-        if result.is_a?(Array)
-          Promise.all(result)
-        else
-          result
-        end
+      begin
+        BatchLoader::Executor.ensure_current
+        blk.call
+      ensure
+        BatchLoader::Executor.clear_current
       end
     end
-
+  
     if max_queries
       result = nil
       expect { result = wrapper.call }.not_to exceed_query_limit(max_queries)