From 9168a19a2d627f33a39a10b5b3fb25d0da7e743f Mon Sep 17 00:00:00 2001
From: Glenn Sarti <glennsarti@users.noreply.github.com>
Date: Wed, 11 Sep 2019 15:29:10 +0800
Subject: [PATCH 1/3] (GH-177) Track dynamic registrations

Previously capability registrations were assumed to only have one possible
reqeuest per method however this isn't true.  Also the request ID is needed in
order to unregister a capability.  This commit;
* Modifies the Language Client to track registrations that are in progress and
  have completed
* Now processes error messages so that registration tracking can fail tracked
  registrations
* Adds tests for these scenarios
---
 lib/puppet-languageserver/language_client.rb  |  48 +++++-
 lib/puppet-languageserver/message_router.rb   |  12 +-
 .../language_client_spec.rb                   | 153 +++++++++++++++++-
 3 files changed, 204 insertions(+), 9 deletions(-)

diff --git a/lib/puppet-languageserver/language_client.rb b/lib/puppet-languageserver/language_client.rb
index 57c87efb..53e5bb4b 100644
--- a/lib/puppet-languageserver/language_client.rb
+++ b/lib/puppet-languageserver/language_client.rb
@@ -4,6 +4,16 @@ module PuppetLanguageServer
   class LanguageClient
     def initialize
       @client_capabilites = {}
+
+      # Internal registry of dynamic registrations and their current state
+      # @registrations[ <[String] method_name>] = [
+      #  {
+      #    :id         => [String] Request ID. Used for de-registration
+      #    :registered => [Boolean] true | false
+      #    :state      => [Enum] :pending | :complete
+      #   }
+      # ]
+      @registrations = {}
     end
 
     def client_capability(*names)
@@ -30,22 +40,52 @@ def parse_lsp_configuration_settings!(settings = [{}])
       # end
     end
 
+    def capability_registrations(method)
+      return [{ :registered => false, :state => :complete }] if @registrations[method].nil?
+      @registrations[method].dup
+    end
+
     def register_capability(message_router, method, options = {})
-      id = SecureRandom.uuid
+      id = new_request_id
 
       PuppetLanguageServer.log_message(:info, "Attempting to dynamically register the #{method} method with id #{id}")
 
+      if @registrations[method] && @registrations[method].select { |i| i[:state] == :pending }.count > 0
+        # The protocol doesn't specify whether this is allowed and is probably per client specific. For the moment we will allow
+        # the registration to be sent but log a message that something may be wrong.
+        PuppetLanguageServer.log_message(:warn, "A dynamic registration for the #{method} method is already in progress")
+      end
+
       params = LSP::RegistrationParams.new.from_h!('registrations' => [])
       params.registrations << LSP::Registration.new.from_h!('id' => id, 'method' => method, 'registerOptions' => options)
+      # Note - Don't put more than one method per request even though you can.  It makes decoding errors much harder!
+
+      @registrations[method] = [] if @registrations[method].nil?
+      @registrations[method] << { :registered => false, :state => :pending, :id => id }
 
       message_router.json_rpc_handler.send_client_request('client/registerCapability', params)
       true
     end
 
-    def parse_register_capability_response!(message_router, _response, original_request)
+    def parse_register_capability_response!(message_router, response, original_request)
       raise 'Response is not from client/registerCapability request' unless original_request['method'] == 'client/registerCapability'
+
+      unless response.key?('result')
+        original_request['params'].registrations.each do |reg|
+          # Mark the registration as completed and failed
+          @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+          @registrations[reg.method__lsp].select { |i| i[:id] == reg.id }.each { |i| i[:registered] = false; i[:state] = :complete } # rubocop:disable Style/Semicolon This is fine
+        end
+        return true
+      end
+
       original_request['params'].registrations.each do |reg|
         PuppetLanguageServer.log_message(:info, "Succesfully dynamically registered the #{reg.method__lsp} method")
+
+        # Mark the registration as completed and succesful
+        @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+        @registrations[reg.method__lsp].select { |i| i[:id] == reg.id }.each { |i| i[:registered] = true; i[:state] = :complete } # rubocop:disable Style/Semicolon This is fine
+
         # If we just registered the workspace/didChangeConfiguration method then
         # also trigger a configuration request to get the initial state
         send_configuration_request(message_router) if reg.method__lsp == 'workspace/didChangeConfiguration'
@@ -56,6 +96,10 @@ def parse_register_capability_response!(message_router, _response, original_requ
 
     private
 
+    def new_request_id
+      SecureRandom.uuid
+    end
+
     def safe_hash_traverse(hash, *names)
       return nil if names.empty?
       item = nil
diff --git a/lib/puppet-languageserver/message_router.rb b/lib/puppet-languageserver/message_router.rb
index d234c556..ce74a0d6 100644
--- a/lib/puppet-languageserver/message_router.rb
+++ b/lib/puppet-languageserver/message_router.rb
@@ -314,15 +314,15 @@ def receive_notification(method, params)
     end
 
     def receive_response(response, original_request)
-      unless response.key?('result')
+      unless receive_response_succesful?(response) # rubocop:disable Style/IfUnlessModifier Too long
         PuppetLanguageServer.log_message(:error, "Response for method '#{original_request['method']}' with id '#{original_request['id']}' failed with #{response['error']}")
-        return
       end
-
+      # Error responses still need to be processed so process messages even if it failed
       case original_request['method']
       when 'client/registerCapability'
         client.parse_register_capability_response!(self, response, original_request)
       when 'workspace/configuration'
+        return unless receive_response_succesful?(response)
         client.parse_lsp_configuration_settings!(response['result'])
       else
         super
@@ -331,6 +331,12 @@ def receive_response(response, original_request)
       PuppetLanguageServer::CrashDump.write_crash_file(e, nil, 'response' => response, 'original_request' => original_request)
       raise
     end
+
+    private
+
+    def receive_response_succesful?(response)
+      response.key?('result')
+    end
   end
 
   class DisabledMessageRouter < BaseMessageRouter
diff --git a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
index 6ab62fa1..9d7bd016 100644
--- a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
+++ b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
@@ -142,6 +142,10 @@
   let(:json_rpc_handler) { MockJSONRPCHandler.new }
   let(:message_router) { MockMessageRouter.new.tap { |i| i.json_rpc_handler = json_rpc_handler } }
 
+  before(:each) do
+    allow(PuppetLanguageServer).to receive(:log_message)
+  end
+
   describe '#client_capability' do
     before(:each) do
       subject.parse_lsp_initialize!(initialize_params)
@@ -180,6 +184,106 @@
     # TODO: Future use.
   end
 
+  describe '#capability_registrations' do
+    let(:method_name) { 'mockMethod' }
+    let(:method_options) { {} }
+    let(:request_id) { 'id001' }
+
+    it 'defaults to false' do
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
+    end
+
+    it 'should track the registration process as it completes succesfully' do
+      req_method_name = nil
+      req_method_params = nil
+      # Remember the registration so we can fake a response later
+      allow(json_rpc_handler).to receive(:send_client_request) do |n, p|
+        req_method_name = n
+        req_method_params = p
+      end
+      # Fake the request id
+      allow(subject).to receive(:new_request_id).and_return(request_id)
+
+      # Should start out not registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
+      # Send as registration request
+      subject.register_capability(message_router, method_name, method_options)
+      # Should show as in progress
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :pending, :id => request_id}])
+      # Mock a valid response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
+      subject.parse_register_capability_response!(message_router, response, original_request)
+      # Should show registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => true, :state => :complete, :id => request_id}])
+    end
+
+    it 'should track the registration process as it fails' do
+      req_method_name = nil
+      req_method_params = nil
+      # Remember the registration so we can fake a response later
+      allow(json_rpc_handler).to receive(:send_client_request) do |n, p|
+        req_method_name = n
+        req_method_params = p
+      end
+      # Fake the request id
+      allow(subject).to receive(:new_request_id).and_return(request_id)
+
+      # Should start out not registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
+      # Send as registration request
+      subject.register_capability(message_router, method_name, method_options)
+      # Should show as in progress
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :pending, :id => request_id}])
+      # Mock an error response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'error' => { 'code' => -1, 'message' => 'mock message' } }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
+      subject.parse_register_capability_response!(message_router, response, original_request)
+      # Should show registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete, :id => request_id}])
+    end
+
+    it 'should preserve the registration state until it is completed' do
+      req_method_name = nil
+      req_method_params = nil
+      # Remember the registration so we can fake a response later
+      allow(json_rpc_handler).to receive(:send_client_request) do |n, p|
+        req_method_name = n
+        req_method_params = p
+      end
+      # Fake the request id
+      request_id2 = 'id002'
+      allow(subject).to receive(:new_request_id).and_return(request_id, request_id2)
+
+      # Should start out not registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
+      # Send as registration request
+      subject.register_capability(message_router, method_name, method_options)
+      # Mock a valid response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
+      subject.parse_register_capability_response!(message_router, response, original_request)
+      # Should show registered
+      expect(subject.capability_registrations(method_name)).to eq([{:registered => true, :state => :complete, :id => request_id}])
+      # Send another registration request
+      subject.register_capability(message_router, method_name, method_options)
+      # Should show as in progress
+      expect(subject.capability_registrations(method_name)).to eq([
+        {:registered => true,  :state => :complete, :id => request_id},
+        {:registered => false, :state => :pending,  :id => request_id2}
+      ])
+      # Mock an error response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'error' => { 'code' => -1, 'message' => 'mock message' } }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
+      subject.parse_register_capability_response!(message_router, response, original_request)
+      # Should still show registered
+      expect(subject.capability_registrations(method_name)).to eq([
+        {:registered => true,  :state => :complete, :id => request_id},
+        {:registered => false, :state => :complete,  :id => request_id2}
+      ])
+    end
+  end
+
   describe '#register_capability' do
     let(:method_name) { 'mockMethod' }
     let(:method_options) { {} }
@@ -198,6 +302,33 @@
       subject.register_capability(message_router, method_name, method_options)
       expect(json_rpc_handler.connection.buffer).to include('"registerOptions":{}')
     end
+
+    it 'should log a message if a registration is already in progress' do
+      allow(json_rpc_handler).to receive(:send_client_request)
+      expect(PuppetLanguageServer).to receive(:log_message).with(:warn, /#{method_name}/)
+
+      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(message_router, method_name, method_options)
+    end
+
+    it 'should not log a message if a previous registration completed' do
+      method_name = nil
+      method_params = nil
+      # Remember the registration so we can fake a response later
+      allow(json_rpc_handler).to receive(:send_client_request) do |n, p|
+        method_name = n
+        method_params = p
+      end
+      expect(PuppetLanguageServer).to_not receive(:log_message).with(:warn, /#{method_name}/)
+      # Send as registration request
+      subject.register_capability(message_router, method_name, method_options)
+      # Mock a valid response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => method_name, 'params' => method_params }
+      subject.parse_register_capability_response!(message_router, response, original_request)
+
+      subject.register_capability(message_router, method_name, method_options)
+    end
   end
 
   describe '#parse_register_capability_response!' do
@@ -237,11 +368,25 @@
         params
       end
 
-      it 'should log the registration' do
-        allow(PuppetLanguageServer).to receive(:log_message)
-        expect(PuppetLanguageServer).to receive(:log_message).with(:info, /validMethod/)
+      context 'that failed' do
+        before(:each) do
+          response.delete('result') if response.key?('result')
+          response['error'] = { 'code' => -1, 'message' => 'mock message' }
+        end
 
-        subject.parse_register_capability_response!(message_router, response, original_request)
+        it 'should not log the registration' do
+          expect(PuppetLanguageServer).to_not receive(:log_message).with(:info, /validMethod/)
+
+          subject.parse_register_capability_response!(message_router, response, original_request)
+        end
+      end
+
+      context 'that succeeded' do
+        it 'should log the registration' do
+          expect(PuppetLanguageServer).to receive(:log_message).with(:info, /validMethod/)
+
+          subject.parse_register_capability_response!(message_router, response, original_request)
+        end
       end
     end
   end

From 9362435993e277fd498c0c966dbd5f678e1b7a33 Mon Sep 17 00:00:00 2001
From: Glenn Sarti <glennsarti@users.noreply.github.com>
Date: Mon, 16 Sep 2019 20:36:13 +0800
Subject: [PATCH 2/3] (GH-177) Refactor message_router as property of Language
 Client

Previously the message router was passed in as a method parameter however as
almost all methods require the message router, this commit changes the
initializer to pass in the message router at object create time.
---
 lib/puppet-languageserver/language_client.rb  | 13 +++--
 lib/puppet-languageserver/message_router.rb   | 10 ++--
 .../language_client_spec.rb                   | 52 +++++++++----------
 .../message_router_spec.rb                    |  6 +--
 4 files changed, 42 insertions(+), 39 deletions(-)

diff --git a/lib/puppet-languageserver/language_client.rb b/lib/puppet-languageserver/language_client.rb
index 53e5bb4b..84cfa95e 100644
--- a/lib/puppet-languageserver/language_client.rb
+++ b/lib/puppet-languageserver/language_client.rb
@@ -2,7 +2,10 @@
 
 module PuppetLanguageServer
   class LanguageClient
-    def initialize
+    attr_reader :message_router
+
+    def initialize(message_router)
+      @message_router = message_router
       @client_capabilites = {}
 
       # Internal registry of dynamic registrations and their current state
@@ -20,7 +23,7 @@ def client_capability(*names)
       safe_hash_traverse(@client_capabilites, *names)
     end
 
-    def send_configuration_request(message_router)
+    def send_configuration_request
       params = LSP::ConfigurationParams.new.from_h!('items' => [])
       params.items << LSP::ConfigurationItem.new.from_h!('section' => 'puppet')
 
@@ -45,7 +48,7 @@ def capability_registrations(method)
       @registrations[method].dup
     end
 
-    def register_capability(message_router, method, options = {})
+    def register_capability(method, options = {})
       id = new_request_id
 
       PuppetLanguageServer.log_message(:info, "Attempting to dynamically register the #{method} method with id #{id}")
@@ -67,7 +70,7 @@ def register_capability(message_router, method, options = {})
       true
     end
 
-    def parse_register_capability_response!(message_router, response, original_request)
+    def parse_register_capability_response!(response, original_request)
       raise 'Response is not from client/registerCapability request' unless original_request['method'] == 'client/registerCapability'
 
       unless response.key?('result')
@@ -88,7 +91,7 @@ def parse_register_capability_response!(message_router, response, original_reque
 
         # If we just registered the workspace/didChangeConfiguration method then
         # also trigger a configuration request to get the initial state
-        send_configuration_request(message_router) if reg.method__lsp == 'workspace/didChangeConfiguration'
+        send_configuration_request if reg.method__lsp == 'workspace/didChangeConfiguration'
       end
 
       true
diff --git a/lib/puppet-languageserver/message_router.rb b/lib/puppet-languageserver/message_router.rb
index ce74a0d6..2f7fff40 100644
--- a/lib/puppet-languageserver/message_router.rb
+++ b/lib/puppet-languageserver/message_router.rb
@@ -34,7 +34,7 @@ class MessageRouter < BaseMessageRouter
     def initialize(options = {})
       super
       @server_options = options.nil? ? {} : options
-      @client = LanguageClient.new
+      @client = LanguageClient.new(self)
     end
 
     def documents
@@ -254,7 +254,7 @@ def receive_notification(method, params)
       when 'initialized'
         PuppetLanguageServer.log_message(:info, 'Client has received initialization')
         if client.client_capability('workspace', 'didChangeConfiguration', 'dynamicRegistration') == true
-          client.register_capability(self, 'workspace/didChangeConfiguration')
+          client.register_capability('workspace/didChangeConfiguration')
         else
           PuppetLanguageServer.log_message(:debug, 'Client does not support didChangeConfiguration dynamic registration. Using push method for configuration change detection.')
         end
@@ -300,7 +300,7 @@ def receive_notification(method, params)
         if params.key?('settings') && params['settings'].nil?
           # This is a notification from a dynamic registration.  Need to send a workspace/configuration
           # request to get the actual configuration
-          client.send_configuration_request(self)
+          client.send_configuration_request
         else
           client.parse_lsp_configuration_settings!(params['settings'])
         end
@@ -314,13 +314,13 @@ def receive_notification(method, params)
     end
 
     def receive_response(response, original_request)
-      unless receive_response_succesful?(response) # rubocop:disable Style/IfUnlessModifier Too long
+      unless receive_response_succesful?(response) # rubocop:disable Style/IfUnlessModifier Line is too long otherwise
         PuppetLanguageServer.log_message(:error, "Response for method '#{original_request['method']}' with id '#{original_request['id']}' failed with #{response['error']}")
       end
       # Error responses still need to be processed so process messages even if it failed
       case original_request['method']
       when 'client/registerCapability'
-        client.parse_register_capability_response!(self, response, original_request)
+        client.parse_register_capability_response!(response, original_request)
       when 'workspace/configuration'
         return unless receive_response_succesful?(response)
         client.parse_lsp_configuration_settings!(response['result'])
diff --git a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
index 9d7bd016..51c8affb 100644
--- a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
+++ b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
@@ -1,8 +1,10 @@
 require 'spec_helper'
 
 describe 'PuppetLanguageServer::LanguageClient' do
+  let(:json_rpc_handler) { MockJSONRPCHandler.new }
+  let(:message_router) { MockMessageRouter.new.tap { |i| i.json_rpc_handler = json_rpc_handler } }
   let(:subject_options) {}
-  let(:subject) { PuppetLanguageServer::LanguageClient.new }
+  let(:subject) { PuppetLanguageServer::LanguageClient.new(message_router) }
   let(:initialize_params) do
     # Example capabilities from VS Code
     {
@@ -139,8 +141,6 @@
       }
     }
   end
-  let(:json_rpc_handler) { MockJSONRPCHandler.new }
-  let(:message_router) { MockMessageRouter.new.tap { |i| i.json_rpc_handler = json_rpc_handler } }
 
   before(:each) do
     allow(PuppetLanguageServer).to receive(:log_message)
@@ -167,11 +167,11 @@
   describe '#send_configuration_request' do
     it 'should send a client request and return true' do
       expect(json_rpc_handler).to receive(:send_client_request).with('workspace/configuration', Object)
-      expect(subject.send_configuration_request(message_router)).to eq(true)
+      expect(subject.send_configuration_request).to eq(true)
     end
 
     it 'should include the puppet settings' do
-      subject.send_configuration_request(message_router)
+      subject.send_configuration_request
       expect(json_rpc_handler.connection.buffer).to include('{"section":"puppet"}')
     end
   end
@@ -207,13 +207,13 @@
       # Should start out not registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
       # Send as registration request
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       # Should show as in progress
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :pending, :id => request_id}])
       # Mock a valid response
       response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
       original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
-      subject.parse_register_capability_response!(message_router, response, original_request)
+      subject.parse_register_capability_response!(response, original_request)
       # Should show registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => true, :state => :complete, :id => request_id}])
     end
@@ -232,13 +232,13 @@
       # Should start out not registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
       # Send as registration request
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       # Should show as in progress
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :pending, :id => request_id}])
       # Mock an error response
       response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'error' => { 'code' => -1, 'message' => 'mock message' } }
       original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
-      subject.parse_register_capability_response!(message_router, response, original_request)
+      subject.parse_register_capability_response!(response, original_request)
       # Should show registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete, :id => request_id}])
     end
@@ -258,15 +258,15 @@
       # Should start out not registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => false, :state => :complete}])
       # Send as registration request
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       # Mock a valid response
       response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
       original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
-      subject.parse_register_capability_response!(message_router, response, original_request)
+      subject.parse_register_capability_response!(response, original_request)
       # Should show registered
       expect(subject.capability_registrations(method_name)).to eq([{:registered => true, :state => :complete, :id => request_id}])
       # Send another registration request
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       # Should show as in progress
       expect(subject.capability_registrations(method_name)).to eq([
         {:registered => true,  :state => :complete, :id => request_id},
@@ -275,7 +275,7 @@
       # Mock an error response
       response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'error' => { 'code' => -1, 'message' => 'mock message' } }
       original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
-      subject.parse_register_capability_response!(message_router, response, original_request)
+      subject.parse_register_capability_response!(response, original_request)
       # Should still show registered
       expect(subject.capability_registrations(method_name)).to eq([
         {:registered => true,  :state => :complete, :id => request_id},
@@ -290,16 +290,16 @@
 
     it 'should send a client request and return true' do
       expect(json_rpc_handler).to receive(:send_client_request).with('client/registerCapability', Object)
-      expect(subject.register_capability(message_router, method_name, method_options)).to eq(true)
+      expect(subject.register_capability(method_name, method_options)).to eq(true)
     end
 
     it 'should include the method to register' do
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       expect(json_rpc_handler.connection.buffer).to include("\"method\":\"#{method_name}\"")
     end
 
     it 'should include the parameters to register' do
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       expect(json_rpc_handler.connection.buffer).to include('"registerOptions":{}')
     end
 
@@ -307,8 +307,8 @@
       allow(json_rpc_handler).to receive(:send_client_request)
       expect(PuppetLanguageServer).to receive(:log_message).with(:warn, /#{method_name}/)
 
-      subject.register_capability(message_router, method_name, method_options)
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
+      subject.register_capability(method_name, method_options)
     end
 
     it 'should not log a message if a previous registration completed' do
@@ -321,13 +321,13 @@
       end
       expect(PuppetLanguageServer).to_not receive(:log_message).with(:warn, /#{method_name}/)
       # Send as registration request
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
       # Mock a valid response
       response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
       original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => method_name, 'params' => method_params }
-      subject.parse_register_capability_response!(message_router, response, original_request)
+      subject.parse_register_capability_response!(response, original_request)
 
-      subject.register_capability(message_router, method_name, method_options)
+      subject.register_capability(method_name, method_options)
     end
   end
 
@@ -342,7 +342,7 @@
       let(:request_params) { {} }
 
       it 'should raise an error if the original request was not a registration' do
-        expect{ subject.parse_register_capability_response!(message_router, response, original_request) }.to raise_error(/client\/registerCapability/)
+        expect{ subject.parse_register_capability_response!(response, original_request) }.to raise_error(/client\/registerCapability/)
       end
     end
 
@@ -355,8 +355,8 @@
       end
 
       it 'should send a configuration request' do
-        expect(subject).to receive(:send_configuration_request).with(message_router)
-        subject.parse_register_capability_response!(message_router, response, original_request)
+        expect(subject).to receive(:send_configuration_request)
+        subject.parse_register_capability_response!(response, original_request)
       end
     end
 
@@ -377,7 +377,7 @@
         it 'should not log the registration' do
           expect(PuppetLanguageServer).to_not receive(:log_message).with(:info, /validMethod/)
 
-          subject.parse_register_capability_response!(message_router, response, original_request)
+          subject.parse_register_capability_response!(response, original_request)
         end
       end
 
@@ -385,7 +385,7 @@
         it 'should log the registration' do
           expect(PuppetLanguageServer).to receive(:log_message).with(:info, /validMethod/)
 
-          subject.parse_register_capability_response!(message_router, response, original_request)
+          subject.parse_register_capability_response!(response, original_request)
         end
       end
     end
diff --git a/spec/languageserver/unit/puppet-languageserver/message_router_spec.rb b/spec/languageserver/unit/puppet-languageserver/message_router_spec.rb
index 34bbf79d..e6aba19f 100644
--- a/spec/languageserver/unit/puppet-languageserver/message_router_spec.rb
+++ b/spec/languageserver/unit/puppet-languageserver/message_router_spec.rb
@@ -858,7 +858,7 @@
         end
 
         it 'should attempt to register workspace/didChangeConfiguration' do
-          expect(subject.client).to receive(:register_capability).with(Object, 'workspace/didChangeConfiguration')
+          expect(subject.client).to receive(:register_capability).with('workspace/didChangeConfiguration')
 
           subject.receive_notification(notification_method, notification_params)
         end
@@ -1037,7 +1037,7 @@
         let(:config_settings) { nil }
 
         it 'should send a configuration request' do
-          expect(subject.client).to receive(:send_configuration_request).with(Object)
+          expect(subject.client).to receive(:send_configuration_request).with(no_args)
 
           subject.receive_notification(notification_method, notification_params)
         end
@@ -1116,7 +1116,7 @@
       let(:request_params) { {} }
 
       it 'should call client.parse_register_capability_response!' do
-        expect(subject.client).to receive(:parse_register_capability_response!).with(Object, response, original_request)
+        expect(subject.client).to receive(:parse_register_capability_response!).with(response, original_request)
 
         subject.receive_response(response, original_request)
       end

From b0e3a6730588fd32b3e36288d312543daa668b68 Mon Sep 17 00:00:00 2001
From: Glenn Sarti <glennsarti@users.noreply.github.com>
Date: Wed, 11 Sep 2019 17:16:46 +0800
Subject: [PATCH 3/3] (GH-177) Add unregistering of capabilities

Previously client capabilities could be dynamically registered.  This commit
adds the logic and track to dynamically unregister a capability;
* Adds the unregister_capability and parse_unregister_capability_response!
  methods
* Adds the client/unregisterCapability notification handler
* Adds tests for unregistering
---
 lib/puppet-languageserver/language_client.rb  |  53 ++++++-
 lib/puppet-languageserver/message_router.rb   |   2 +
 .../language_client_spec.rb                   | 146 ++++++++++++++++++
 3 files changed, 199 insertions(+), 2 deletions(-)

diff --git a/lib/puppet-languageserver/language_client.rb b/lib/puppet-languageserver/language_client.rb
index 84cfa95e..22351292 100644
--- a/lib/puppet-languageserver/language_client.rb
+++ b/lib/puppet-languageserver/language_client.rb
@@ -44,7 +44,7 @@ def parse_lsp_configuration_settings!(settings = [{}])
     end
 
     def capability_registrations(method)
-      return [{ :registered => false, :state => :complete }] if @registrations[method].nil?
+      return [{ :registered => false, :state => :complete }] if @registrations[method].nil? || @registrations[method].empty?
       @registrations[method].dup
     end
 
@@ -56,7 +56,7 @@ def register_capability(method, options = {})
       if @registrations[method] && @registrations[method].select { |i| i[:state] == :pending }.count > 0
         # The protocol doesn't specify whether this is allowed and is probably per client specific. For the moment we will allow
         # the registration to be sent but log a message that something may be wrong.
-        PuppetLanguageServer.log_message(:warn, "A dynamic registration for the #{method} method is already in progress")
+        PuppetLanguageServer.log_message(:warn, "A dynamic registration/deregistration for the #{method} method is already in progress")
       end
 
       params = LSP::RegistrationParams.new.from_h!('registrations' => [])
@@ -70,6 +70,31 @@ def register_capability(method, options = {})
       true
     end
 
+    def unregister_capability(method)
+      if @registrations[method].nil?
+        PuppetLanguageServer.log_message(:debug, "No registrations to deregister for the #{method}")
+        return true
+      end
+
+      params = LSP::UnregistrationParams.new.from_h!('unregisterations' => [])
+      @registrations[method].each do |reg|
+        next if reg[:id].nil?
+        PuppetLanguageServer.log_message(:warn, "A dynamic registration/deregistration for the #{method} method, with id #{reg[:id]} is already in progress") if reg[:state] == :pending
+        # Ignore registrations that don't need to be unregistered
+        next if reg[:state] == :complete && !reg[:registered]
+        params.unregisterations << LSP::Unregistration.new.from_h!('id' => reg[:id], 'method' => method)
+        reg[:state] = :pending
+      end
+
+      if params.unregisterations.count.zero?
+        PuppetLanguageServer.log_message(:debug, "Nothing to deregister for the #{method} method")
+        return true
+      end
+
+      message_router.json_rpc_handler.send_client_request('client/unregisterCapability', params)
+      true
+    end
+
     def parse_register_capability_response!(response, original_request)
       raise 'Response is not from client/registerCapability request' unless original_request['method'] == 'client/registerCapability'
 
@@ -97,6 +122,30 @@ def parse_register_capability_response!(response, original_request)
       true
     end
 
+    def parse_unregister_capability_response!(response, original_request)
+      raise 'Response is not from client/unregisterCapability request' unless original_request['method'] == 'client/unregisterCapability'
+
+      unless response.key?('result')
+        original_request['params'].unregisterations.each do |reg|
+          # Mark the registration as completed and failed
+          @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+          @registrations[reg.method__lsp].select { |i| i[:id] == reg.id && i[:registered] }.each { |i| i[:state] = :complete }
+          @registrations[reg.method__lsp].delete_if { |i| i[:id] == reg.id && !i[:registered] }
+        end
+        return true
+      end
+
+      original_request['params'].unregisterations.each do |reg|
+        PuppetLanguageServer.log_message(:info, "Succesfully dynamically unregistered the #{reg.method__lsp} method")
+
+        # Remove registrations
+        @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+        @registrations[reg.method__lsp].delete_if { |i| i[:id] == reg.id }
+      end
+
+      true
+    end
+
     private
 
     def new_request_id
diff --git a/lib/puppet-languageserver/message_router.rb b/lib/puppet-languageserver/message_router.rb
index 2f7fff40..f6cd6b9c 100644
--- a/lib/puppet-languageserver/message_router.rb
+++ b/lib/puppet-languageserver/message_router.rb
@@ -321,6 +321,8 @@ def receive_response(response, original_request)
       case original_request['method']
       when 'client/registerCapability'
         client.parse_register_capability_response!(response, original_request)
+      when 'client/unregisterCapability'
+        client.parse_unregister_capability_response!(response, original_request)
       when 'workspace/configuration'
         return unless receive_response_succesful?(response)
         client.parse_lsp_configuration_settings!(response['result'])
diff --git a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
index 51c8affb..22a48eca 100644
--- a/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
+++ b/spec/languageserver/unit/puppet-languageserver/language_client_spec.rb
@@ -331,6 +331,72 @@
     end
   end
 
+  describe '#unregister_capability' do
+    let(:method_name) { 'mockMethod' }
+
+    before(:each) do
+      # Mock an already succesful registration
+      subject.instance_variable_set(:@registrations, {
+        method_name => [{ :id => 'id001', :state => :complete, :registered => true }]
+      })
+    end
+
+    it 'should send a client request and return true' do
+      expect(json_rpc_handler).to receive(:send_client_request).with('client/unregisterCapability', Object)
+      expect(subject.unregister_capability(method_name)).to eq(true)
+    end
+
+    it 'should include the method to register' do
+      subject.unregister_capability(method_name)
+      expect(json_rpc_handler.connection.buffer).to include("\"method\":\"#{method_name}\"")
+    end
+
+    it 'should log a message if a registration is already in progress' do
+      allow(json_rpc_handler).to receive(:send_client_request)
+      expect(PuppetLanguageServer).to receive(:log_message).with(:warn, /#{method_name}/)
+
+      subject.unregister_capability(method_name)
+      subject.unregister_capability(method_name)
+    end
+
+    it 'should not log a message if a previous registration completed' do
+      req_method_name = nil
+      req_method_params = nil
+      # Remember the registration so we can fake a response later
+      allow(json_rpc_handler).to receive(:send_client_request) do |n, p|
+        req_method_name = n
+        req_method_params = p
+      end
+
+      expect(PuppetLanguageServer).to_not receive(:log_message).with(:warn, /#{method_name}/)
+      # Send as registration request
+      subject.unregister_capability(method_name)
+      # Mock a valid response
+      response = { 'jsonrpc'=>'2.0', 'id'=> 0, 'result' => nil }
+      original_request = { 'jsonrpc'=>'2.0', 'id' => 0, 'method' => req_method_name, 'params' => req_method_params }
+
+      subject.parse_unregister_capability_response!(response, original_request)
+
+      subject.unregister_capability(method_name)
+    end
+
+    it 'should not deregister methods that have not been registerd' do
+      expect(json_rpc_handler).to_not receive(:send_client_request)
+
+      subject.unregister_capability('unknown')
+    end
+
+    it 'should not deregister methods that are no longer registerd' do
+      expect(json_rpc_handler).to_not receive(:send_client_request)
+
+      subject.instance_variable_set(:@registrations, {
+        method_name => [{ :id => 'id001', :state => :complete, :registered => false }]
+      })
+
+      subject.unregister_capability(method_name)
+    end
+  end
+
   describe '#parse_register_capability_response!' do
     let(:request_id) { 0 }
     let(:response_result) { nil }
@@ -390,4 +456,84 @@
       end
     end
   end
+
+  describe '#parse_unregister_capability_response!' do
+    let(:request_id) { 0 }
+    let(:response_result) { nil }
+    let(:response) { {'jsonrpc'=>'2.0', 'id'=> request_id, 'result' => response_result } }
+    let(:original_request) { {'jsonrpc'=>'2.0', 'id'=> request_id, 'method' => request_method, 'params' => request_params} }
+    let(:method_name) { 'validMethod' }
+    let(:initial_registration) { true }
+
+    before(:each) do
+      # Mock an already succesful registration
+      subject.instance_variable_set(:@registrations, {
+        method_name => [{ :id => 'id001', :state => :complete, :registered => initial_registration }]
+      })
+    end
+
+    context 'Given an original request that is not an unregistration' do
+      let(:request_method) { 'mockMethod' }
+      let(:request_params) { {} }
+
+      it 'should raise an error if the original request was not a registration' do
+        expect{ subject.parse_unregister_capability_response!(response, original_request) }.to raise_error(/client\/unregisterCapability/)
+      end
+    end
+
+    context 'Given a valid original request' do
+      let(:request_method) { 'client/unregisterCapability' }
+      let(:request_params) do
+        params = LSP::UnregistrationParams.new.from_h!('unregisterations' => [])
+        params.unregisterations << LSP::Unregistration.new.from_h!('id' => 'id001', 'method' => method_name)
+        params
+      end
+
+      before(:each) do
+        # Mimic an unregistration that is in progress
+        subject.instance_variable_set(:@registrations, {
+          method_name => [{ :id => 'id001', :state => :pending, :registered => initial_registration }]
+        })
+      end
+
+      context 'that failed' do
+        before(:each) do
+          response.delete('result') if response.key?('result')
+          response['error'] = { 'code' => -1, 'message' => 'mock message' }
+        end
+
+        context 'and was previously registered' do
+          it 'should retain that it is registered' do
+            subject.parse_unregister_capability_response!(response, original_request)
+
+            expect(subject.capability_registrations(method_name)).to eq([{:id=>"id001", :registered=>true, :state=>:complete}])
+          end
+        end
+
+        context 'and was not previously registered' do
+          let(:initial_registration) { false }
+
+          it 'should no longer be in the registration list' do
+            subject.parse_unregister_capability_response!(response, original_request)
+
+            expect(subject.capability_registrations(method_name)).to eq([{ :registered => false, :state => :complete }])
+          end
+        end
+      end
+
+      context 'that succeeded' do
+        it 'should log the registration' do
+          expect(PuppetLanguageServer).to receive(:log_message).with(:info, /validMethod/)
+
+          subject.parse_unregister_capability_response!(response, original_request)
+        end
+
+        it 'should no longer be in the registration list' do
+          subject.parse_unregister_capability_response!(response, original_request)
+
+          expect(subject.capability_registrations(method_name)).to eq([{ :registered => false, :state => :complete }])
+        end
+      end
+    end
+  end
 end