From 934729fcf05baf6e6074ef85b9be2784dc1b6f17 Mon Sep 17 00:00:00 2001
From: Andy Waite <andy.waite@shopify.com>
Date: Thu, 19 Dec 2024 15:19:53 -0500
Subject: [PATCH] Add test helper methods from Rails add-on (#2991)

* Add test helper methods from Rails add-on
---
 .rubocop.yml                |  1 +
 lib/ruby_lsp/test_helper.rb | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/.rubocop.yml b/.rubocop.yml
index 2ba910dcb..4fc414f9d 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -60,6 +60,7 @@ Sorbet/StrictSigil:
     - "lib/ruby_indexer/lib/ruby_indexer/prefix_tree.rb"
     - "lib/ruby_lsp/load_sorbet.rb"
     - "lib/ruby_lsp/scripts/compose_bundle.rb"
+    - "lib/ruby_lsp/test_helper.rb"
 
 Layout/ClassStructure:
   Enabled: true
diff --git a/lib/ruby_lsp/test_helper.rb b/lib/ruby_lsp/test_helper.rb
index 98ecf6bf1..d905dd3e4 100644
--- a/lib/ruby_lsp/test_helper.rb
+++ b/lib/ruby_lsp/test_helper.rb
@@ -1,11 +1,16 @@
-# typed: strict
+# typed: true
 # frozen_string_literal: true
 
 # NOTE: This module is intended to be used by addons for writing their own tests, so keep that in mind if changing.
 
 module RubyLsp
   module TestHelper
+    class TestError < StandardError; end
+
     extend T::Sig
+    extend T::Helpers
+
+    requires_ancestor { Kernel }
 
     sig do
       type_parameters(:T)
@@ -52,5 +57,33 @@ def with_server(source = nil, uri = Kernel.URI("file:///fake.rb"), stub_no_typec
         server.run_shutdown
       end
     end
+
+    sig { params(server: RubyLsp::Server).returns(RubyLsp::Result) }
+    def pop_result(server)
+      result = server.pop_response
+      result = server.pop_response until result.is_a?(RubyLsp::Result) || result.is_a?(RubyLsp::Error)
+
+      if result.is_a?(RubyLsp::Error)
+        raise TestError, "Failed to execute request #{result.message}"
+      else
+        result
+      end
+    end
+
+    def pop_log_notification(message_queue, type)
+      log = message_queue.pop
+      return log if log.params.type == type
+
+      log = message_queue.pop until log.params.type == type
+      log
+    end
+
+    def pop_message(outgoing_queue, &block)
+      message = outgoing_queue.pop
+      return message if block.call(message)
+
+      message = outgoing_queue.pop until block.call(message)
+      message
+    end
   end
 end