1 files changed, 29 insertions, 4 deletions
diff --git a/test/test_find_references.vader b/test/test_find_references.vader
index 150e0471..ecced068 100644
--- a/test/test_find_references.vader
+++ b/test/test_find_references.vader
@@ -3,20 +3,20 @@ Before:
   call ale#test#SetFilename('dummy.txt')
 
   let g:old_filename = expand('%:p')
-  let g:Callback = 0
+  let g:Callback = ''
   let g:expr_list = []
   let g:message_list = []
   let g:preview_called = 0
   let g:item_list = []
+  let g:capability_checked = ''
+  let g:WaitCallback = v:null
 
   runtime autoload/ale/linter.vim
   runtime autoload/ale/lsp.vim
   runtime autoload/ale/util.vim
   runtime autoload/ale/preview.vim
 
-  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}
@@ -29,6 +29,15 @@ Before:
     \}
   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
+
   function! ale#lsp#Send(conn_id, message, root) abort
     call add(g:message_list, a:message)
 
@@ -50,6 +59,8 @@ After:
   call ale#test#RestoreDirectory()
   call ale#linter#Reset()
 
+  unlet! g:capability_checked
+  unlet! g:WaitCallback
   unlet! g:old_filename
   unlet! g:Callback
   unlet! g:message_list
@@ -152,6 +163,13 @@ Execute(tsserver reference requests should be sent):
 
   ALEFindReferences
 
+  " We shouldn't register the callback yet.
+  AssertEqual '''''', string(g:Callback)
+
+  AssertEqual type(function('type')), type(g:WaitCallback)
+  AssertEqual 'references', g:capability_checked
+  call call(g:WaitCallback, [347, '/foo/bar'])
+
   AssertEqual
   \ 'function(''ale#references#HandleTSServerResponse'')',
   \ string(g:Callback)
@@ -226,6 +244,13 @@ Execute(LSP reference requests should be sent):
 
   ALEFindReferences
 
+  " We shouldn't register the callback yet.
+  AssertEqual '''''', string(g:Callback)
+
+  AssertEqual type(function('type')), type(g:WaitCallback)
+  AssertEqual 'references', g:capability_checked
+  call call(g:WaitCallback, [347, '/foo/bar'])
+
   AssertEqual
   \ 'function(''ale#references#HandleLSPResponse'')',
   \ string(g:Callback)