1 files changed, 28 insertions, 3 deletions
diff --git a/test/completion/test_lsp_completion_messages.vader b/test/completion/test_lsp_completion_messages.vader
index 00a174dc..30a9c8e1 100644
--- a/test/completion/test_lsp_completion_messages.vader
+++ b/test/completion/test_lsp_completion_messages.vader
@@ -13,11 +13,11 @@ Before:
   runtime autoload/ale/lsp.vim
 
   let g:message_list = []
+  let g:capability_checked = ''
   let g:Callback = ''
+  let g:WaitCallback = v:null
 
-  function! ale#lsp_linter#StartLSP(buffer, linter, callback) abort
-    let g:Callback = a:callback
-
+  function! ale#lsp_linter#StartLSP(buffer, linter) abort
     let l:conn = ale#lsp#NewConnection({})
     let l:conn.id = 347
     let l:conn.open_documents = {a:buffer : -1}
@@ -35,6 +35,15 @@ Before:
     return 'i'
   endfunction
 
+  function! ale#lsp#WaitForCapability(conn_id, project_root, capability, callback) abort
+    let g:capability_checked = a:capability
+    let g:WaitCallback = a:callback
+  endfunction
+
+  function! ale#lsp#RegisterCallback(conn_id, callback) abort
+    let g:Callback = a:callback
+  endfunction
+
   " Replace the Send function for LSP, so we can monitor calls to it.
   function! ale#lsp#Send(conn_id, message, ...) abort
     call add(g:message_list, a:message)
@@ -44,6 +53,8 @@ After:
   Restore
 
   unlet! g:message_list
+  unlet! g:capability_checked
+  unlet! g:WaitCallback
   unlet! g:Callback
   unlet! b:ale_old_omnifunc
   unlet! b:ale_old_completopt
@@ -84,6 +95,13 @@ Execute(The right message should be sent for the initial tsserver request):
 
   call ale#completion#GetCompletions()
 
+  " We shouldn't register the callback yet.
+  AssertEqual '''''', string(g:Callback)
+
+  AssertEqual type(function('type')), type(g:WaitCallback)
+  AssertEqual 'completion', g:capability_checked
+  call call(g:WaitCallback, [347, '/foo/bar'])
+
   " We should send the right callback.
   AssertEqual
   \ 'function(''ale#completion#HandleTSServerResponse'')',
@@ -164,6 +182,13 @@ Execute(The right message should be sent for the initial LSP request):
 
   call ale#completion#GetCompletions()
 
+  " We shouldn't register the callback yet.
+  AssertEqual '''''', string(g:Callback)
+
+  AssertEqual type(function('type')), type(g:WaitCallback)
+  AssertEqual 'completion', g:capability_checked
+  call call(g:WaitCallback, [347, '/foo/bar'])
+
   " We should send the right callback.
   AssertEqual
   \ 'function(''ale#completion#HandleLSPResponse'')',