From 12955db1deb40739507aca02c069d36b718bdc84 Mon Sep 17 00:00:00 2001
From: Siddarth Challa <sc5455@nyu.edu>
Date: Wed, 11 Dec 2024 16:36:33 -0500
Subject: [PATCH] Bug fixes + finish writing unit tests

---
 lib/ruby_lsp/ruby_lsp_rails/completion.rb | 25 ++++--
 test/ruby_lsp_rails/completion_test.rb    | 94 +++++++++++++++++++++++
 2 files changed, 113 insertions(+), 6 deletions(-)
 create mode 100644 test/ruby_lsp_rails/completion_test.rb

diff --git a/lib/ruby_lsp/ruby_lsp_rails/completion.rb b/lib/ruby_lsp/ruby_lsp_rails/completion.rb
index eb9629e3..eb4f52e9 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/completion.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/completion.rb
@@ -10,9 +10,7 @@ class Completion
       sig do
         override.params(
           client: RunnerClient,
-          response_builder: ResponseBuilders::CollectionResponseBuilder[T.any(
-            Interface::Location, Interface::LocationLink
-          )],
+          response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
           node_context: NodeContext,
           dispatcher: Prism::Dispatcher,
           uri: URI::Generic,
@@ -30,7 +28,7 @@ def initialize(client, response_builder, node_context, dispatcher, uri)
 
       sig { params(node: Prism::CallNode).void }
       def on_call_node_enter(node)
-        if @node_context&.call_node&.name == :where
+        if node.name == :where
           handle_active_record_where_completions(node)
         end
       end
@@ -39,15 +37,30 @@ def on_call_node_enter(node)
 
       sig { params(node: Prism::CallNode).void }
       def handle_active_record_where_completions(node)
-        resolved_class = @client.model(T.must(@node_context.call_node).receiver&.name)
+        receiver = node.receiver
+        return unless receiver
+
+        resolved_class = @client.model(receiver.name.to_s)
         return if resolved_class.nil?
 
+        arguments_node = node.arguments
+        existing_args = {}
+        if arguments_node
+          arguments_node.arguments.each do |arg|
+            arg.elements.each do |a|
+              existing_args[a.key.unescaped] = true
+            end
+          end
+        end
+
         resolved_class[:columns].each do |column|
+          next if existing_args[column[0]]
+
           @response_builder << Interface::CompletionItem.new(
             label: column[0],
             filter_text: column[0],
             label_details: Interface::CompletionItemLabelDetails.new(
-              description: "Filter #{T.must(@node_context.call_node).receiver.name} records by #{column[0]}",
+              description: "Filter #{node.receiver.name} records by #{column[0]}",
             ),
             text_edit: Interface::TextEdit.new(range: 0, new_text: "#{column[0]}: "),
             kind: Constant::CompletionItemKind::FIELD,
diff --git a/test/ruby_lsp_rails/completion_test.rb b/test/ruby_lsp_rails/completion_test.rb
new file mode 100644
index 00000000..f5b8787a
--- /dev/null
+++ b/test/ruby_lsp_rails/completion_test.rb
@@ -0,0 +1,94 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RubyLsp
+  module Rails
+    class CompletionTest < ActiveSupport::TestCase
+      test "recognizes Active Record .where call on an Active Record model when cursor is on (" do
+        response = generate_completions_for_source(<<~RUBY, { line: 2, character: 10 })
+          # typed: false
+
+          User.where(
+        RUBY
+
+        columns = [
+          "id",
+          "first_name",
+          "last_name",
+          "age",
+          "created_at",
+          "updated_at",
+          "country_id",
+          "active",
+        ]
+        assert_equal(columns.size, response.size)
+
+        columns.each_with_index do |column, i|
+          assert_equal(column, response[i].label)
+          assert_equal(column, response[i].filter_text)
+        end
+      end
+
+      test "Does not suggest column if it already exists within .where as an arg and parantheses are not closed" do
+        response = generate_completions_for_source(<<~RUBY, { line: 2, character: 28 })
+          # typed: false
+
+          User.where(id:, first_name:, 
+        RUBY
+
+        columns = [ "last_name", "age", "created_at", "updated_at", "country_id", "active" ]
+        assert_equal(columns.size, response.size)
+
+        columns.each_with_index do |column, i|
+          assert_equal(column, response[i].label)
+          assert_equal(column, response[i].filter_text)
+        end
+      end
+
+      test "Does not suggest column if it already exists within .where as an arg and parantheses are closed" do
+        response = generate_completions_for_source(<<~RUBY, { line: 2, character: 28 })
+          # typed: false
+
+          User.where(id:, first_name:, )
+        RUBY
+
+        columns = [ "last_name", "age", "created_at", "updated_at", "country_id", "active" ]
+        assert_equal(columns.size, response.size)
+
+        columns.each_with_index do |column, i|
+          assert_equal(column, response[i].label)
+          assert_equal(column, response[i].filter_text)
+        end
+      end
+
+      test "Does not provide suggestions when typing argument values" do
+        response = generate_completions_for_source(<<~RUBY, { line: 1, character: 14})
+          # typed: false
+          User.where(id: 
+        RUBY
+
+        assert_equal(0, response.size)
+      end
+
+
+      private
+
+      def generate_completions_for_source(source, position)
+        with_server(source) do |server, uri|
+          sleep(0.1) while RubyLsp::Addon.addons.first.instance_variable_get(:@rails_runner_client).is_a?(NullClient)
+
+          server.process_message(
+            id: 1,
+            method: "textDocument/completion",
+            params: { textDocument: { uri: uri }, position: position },
+          )
+
+          result = pop_result(server)
+          result.response
+        end
+      end
+    end
+  end
+end